ran-proc/rpcore/cptable.go

package rpcore

import (
	"fmt"
	memdb "github.com/hashicorp/go-memdb"
)

const CPTABLE_NAME = "cpt"

type ChildProcTable struct {
	db *memdb.MemDB
}

func NewChildProcTable() *ChildProcTable {
	scheme := &memdb.DBSchema{
		Tables: map[string]*memdb.TableSchema{
			CPTABLE_NAME: &memdb.TableSchema{
				Name: CPTABLE_NAME,
				Indexes: map[string]*memdb.IndexSchema{
					"id": &memdb.IndexSchema{
						Name:         "id",
						AllowMissing: false,
						Unique:       true,
						Indexer:      &memdb.IntFieldIndex{Field: "ID"},
					},
					"name": &memdb.IndexSchema{
						Name:         "name",
						AllowMissing: false,
						Unique:       false,
						Indexer:      &memdb.StringFieldIndex{Field: "Name"},
					},
					"is_daemon": &memdb.IndexSchema{
						Name:         "is_daemon",
						AllowMissing: false,
						Unique:       false,
						Indexer:      &memdb.BoolFieldIndex{Field: "IsDaemon"},
					},
					"status": &memdb.IndexSchema{
						Name:         "status",
						AllowMissing: false,
						Unique:       false,
						Indexer:      &memdb.IntFieldIndex{Field: "Status"},
					},
					"pid": &memdb.IndexSchema{
						Name:         "pid",
						AllowMissing: false,
						Unique:       false,
						Indexer:      &memdb.IntFieldIndex{Field: "ProcessID"},
					},
				},
			},
		},
	}
	mdb, err := memdb.NewMemDB(scheme)
	if err != nil {
		panic(fmt.Errorf("failed init ChildProcTable: memdb error: %s", err))
	}
	t := &ChildProcTable{
		db: mdb,
	}
	return t
}

func (c *ChildProcTable) Add(item *ChildProcTableItem) error {
	txn := c.db.Txn(true)
	err := txn.Insert(CPTABLE_NAME, item)
	if err != nil {
		return fmt.Errorf("add cpitem failed: set: %s", err)
	}
	txn.Commit()
	return nil
}

func (c *ChildProcTable) GetByID(id int64) (*ChildProcTableItem, error) {
	txn := c.db.Txn(false)
	defer txn.Abort()
	raw, err := txn.First(CPTABLE_NAME, "id", id)
	if err != nil {
		return nil, fmt.Errorf("get cpitem failed: %s", err)
	}
	if raw == nil {
		return nil, fmt.Errorf("get cpitem failed: no such item")
	}
	return raw.(*ChildProcTableItem), nil
}

func (c *ChildProcTable) GetByName(name string) ([]*ChildProcTableItem, error) {
	txn := c.db.Txn(false)
	defer txn.Abort()
	it, err := txn.Get(CPTABLE_NAME, "name", name)
	if err != nil {
		return nil, fmt.Errorf("get cpitem failed: %s", err)
	}
	out := make([]*ChildProcTableItem, 0)
	for obj := it.Next(); obj != nil; obj = it.Next() {
		out = append(out, obj.(*ChildProcTableItem))
	}
	return out, nil
}

func (c *ChildProcTable) Delete(item *ChildProcTableItem) error {
	txn := c.db.Txn(true)
	err := txn.Delete(CPTABLE_NAME, item)
	if err != nil {
		return fmt.Errorf("delete cpitem failed: %s", err)
	}
	txn.Commit()
	return nil
}

func (c *ChildProcTable) GetByPID(pid int) ([]*ChildProcTableItem, error) {
	txn := c.db.Txn(false)
	defer txn.Abort()
	it, err := txn.Get(CPTABLE_NAME, "pid", pid)
	if err != nil {
		return nil, fmt.Errorf("get cpitem failed: %s", err)
	}
	out := make([]*ChildProcTableItem, 0)
	for obj := it.Next(); obj != nil; obj = it.Next() {
		out = append(out, obj.(*ChildProcTableItem))
	}
	return out, nil
}

func (c *ChildProcTable) ListAll() ([]*ChildProcTableItem, error) {
	txn := c.db.Txn(false)
	defer txn.Abort()
	it, err := txn.Get(CPTABLE_NAME, "id")
	if err != nil {
		return nil, fmt.Errorf("get items failed: %s", err)
	}
	out := make([]*ChildProcTableItem, 0)
	for obj := it.Next(); obj != nil; obj = it.Next() {
		out = append(out, obj.(*ChildProcTableItem))
	}
	return out, nil
}

func (c *ChildProcTable) DoQuery(index string, args ...interface{}) ([]*ChildProcTableItem, error) {
	txn := c.db.Txn(false)
	defer txn.Abort()
	it, err := txn.Get(CPTABLE_NAME, index, args...)
	if err != nil {
		return nil, fmt.Errorf("query items failed: %s", err)
	}
	out := make([]*ChildProcTableItem, 0)
	for obj := it.Next(); obj != nil; obj = it.Next() {
		out = append(out, obj.(*ChildProcTableItem))
	}
	return out, nil
}