db.go

package main

import (
    "encoding/json"
    "fmt"
)

// DB wraps the KV store to provide a Relational Table interface
type DB struct {
    Path string
    kv   *KV
}

// DBTX represents a high-level Relational database transaction.
type DBTX struct {
    kv KVTX
    db *DB
}

// Begin starts a new relational transaction
func (db *DB) Begin(tx *DBTX) {
    tx.db = db
    db.kv.Begin(&tx.kv)
}

// Commit attempts to save all transaction changes to disk atomically
func (db *DB) Commit(tx *DBTX) error {
    return db.kv.Commit(&tx.kv)
}

// Abort drops all pending changes in this transaction
func (db *DB) Abort(tx *DBTX) {
    db.kv.Abort(&tx.kv)
}

func checkRecord(tdef *TableDef, rec Record, n int) ([]Value, error) {
    values := make([]Value, n)
    for i := 0; i < n; i++ {
        val := rec.Get(rec.Cols[i])
        if val == nil {
            return nil, fmt.Errorf("missing column: %s", rec.Cols[i])
        }
        if val.Type != tdef.Types[i] {
            return nil, fmt.Errorf("type mismatch for column: %s", rec.Cols[i])
        }
        values[i] = *val
    }
    return values, nil
}

func buildIndexKey(tdef *TableDef, indexNum int, rec Record) ([]byte, error) {
    indexCols := tdef.Indexes[indexNum]
    vals := make([]Value, len(indexCols))
    
    for i, col := range indexCols {
        val := rec.Get(col)
        if val == nil {
            return nil, fmt.Errorf("missing indexed column: %s", col)
        }
        vals[i] = *val
    }

    if indexNum > 0 {
        pkCols := tdef.Indexes[0]
        for _, col := range pkCols {
            val := rec.Get(col)
            if val == nil {
                return nil, fmt.Errorf("missing primary key column for secondary index: %s", col)
            }
            vals = append(vals, *val)
        }
    }
    
    return encodeKey(nil, tdef.Prefixes[indexNum], vals), nil
}

func dbGet(tx *DBTX, tdef *TableDef, rec *Record) (bool, error) {
    key, err := buildIndexKey(tdef, 0, *rec)
    if err != nil {
        return false, err
    }

    val, ok := tx.kv.Get(key)
    if !ok {
        return false, nil
    }

    pkLen := len(tdef.Indexes[0])
    values := make([]Value, len(tdef.Cols))

    for i := 0; i < pkLen; i++ {
        values[i].Type = tdef.Types[i]
    }
    decodeValues(key[4:], values[:pkLen])

    for i := pkLen; i < len(tdef.Cols); i++ {
        values[i].Type = tdef.Types[i]
    }
    decodeValues(val, values[pkLen:])

    rec.Cols = tdef.Cols
    rec.Vals = values
    return true, nil
}

func getTableDef(tx *DBTX, name string) *TableDef {
    if name == "@table" {
        return TDEF_TABLE
    }
    if name == "@meta" {
        return TDEF_META
    }

    rec := (&Record{}).AddStr("name", []byte(name))
    ok, err := dbGet(tx, TDEF_TABLE, rec)
    if err != nil || !ok {
        return nil
    }

    tdef := &TableDef{}
    err = json.Unmarshal(rec.Get("def").Str, tdef)
    if err != nil {
        panic(err)
    }
    return tdef
}

func (tx *DBTX) Upsert(table string, rec Record) (bool, error) {
    tdef := getTableDef(tx, table)
    if tdef == nil {
        return false, fmt.Errorf("table not found: %s", table)
    }

    var oldRec Record
    for _, col := range tdef.Indexes[0] { 
        oldRec.Cols = append(oldRec.Cols, col)
        oldRec.Vals = append(oldRec.Vals, *rec.Get(col))
    }

    exists, _ := dbGet(tx, tdef, &oldRec)
    if exists {
        for i := 1; i < len(tdef.Indexes); i++ {
            oldKey, _ := buildIndexKey(tdef, i, oldRec)
            tx.kv.Del(oldKey) // I/O happens in memory until Commit!
        }
    }

    for i := 1; i < len(tdef.Indexes); i++ {
        newKey, _ := buildIndexKey(tdef, i, rec)
        tx.kv.Set(newKey, []byte{}) 
    }

    pkKey, _ := buildIndexKey(tdef, 0, rec)

    pkLen := len(tdef.Indexes[0])
    valVals := make([]Value, len(tdef.Cols)-pkLen)
    for i := pkLen; i < len(tdef.Cols); i++ {
        valVals[i-pkLen] = *rec.Get(tdef.Cols[i])
    }
    pkVal := encodeValues(nil, valVals)

    tx.kv.Set(pkKey, pkVal) // Error return removed because it's delayed to Commit
    return true, nil
}

func (tx *DBTX) Get(table string, rec *Record) (bool, error) {
    tdef := getTableDef(tx, table)
    if tdef == nil {
        return false, fmt.Errorf("table not found: %s", table)
    }
    return dbGet(tx, tdef, rec)
}

func (tx *DBTX) Delete(table string, rec Record) (bool, error) {
    tdef := getTableDef(tx, table)
    if tdef == nil {
        return false, fmt.Errorf("table not found: %s", table)
    }

    var oldRec Record
    for _, col := range tdef.Indexes[0] {
        oldRec.Cols = append(oldRec.Cols, col)
        oldRec.Vals = append(oldRec.Vals, *rec.Get(col))
    }

    exists, _ := dbGet(tx, tdef, &oldRec)
    if !exists {
        return false, nil 
    }

    for i := 1; i < len(tdef.Indexes); i++ {
        oldKey, _ := buildIndexKey(tdef, i, oldRec)
        tx.kv.Del(oldKey)
    }

    pkKey, _ := buildIndexKey(tdef, 0, oldRec)
    tx.kv.Del(pkKey) // Error return removed
    return true, nil
}

func (tx *DBTX) Insert(table string, rec Record) (bool, error) {
    exists, err := tx.Get(table, &rec)
    if err != nil {
        return false, err
    }
    if exists {
        return false, fmt.Errorf("duplicate key error: row already exists")
    }
    return tx.Upsert(table, rec)
}

func (tx *DBTX) Update(table string, rec Record) (bool, error) {
    exists, err := tx.Get(table, &rec)
    if err != nil {
        return false, err
    }
    if !exists {
        return false, fmt.Errorf("update error: row does not exist")
    }
    return tx.Upsert(table, rec)
}

func (tx *DBTX) TableNew(tdef *TableDef) error {
    if getTableDef(tx, tdef.Name) != nil {
        return fmt.Errorf("table already exists: %s", tdef.Name)
    }

    metaKey := (&Record{}).AddStr("key", []byte("next_prefix"))
    ok, err := dbGet(tx, TDEF_META, metaKey)
    if err != nil {
        return err
    }

    var nextPrefix uint32 = 3
    if ok {
        nextPrefix = uint32(metaKey.Get("val").I64)
    }

    tdef.Prefixes = make([]uint32, len(tdef.Indexes))
    for i := range tdef.Indexes {
        tdef.Prefixes[i] = nextPrefix + uint32(i)
    }

    defBytes, err := json.Marshal(tdef)
    if err != nil {
        return err
    }

    tableRec := (&Record{}).
        AddStr("name", []byte(tdef.Name)).
        AddStr("def", defBytes)

    if _, err := tx.Upsert("@table", *tableRec); err != nil {
        return err
    }

    metaKey.Cols = nil
    metaKey.Vals = nil
    metaKey.AddStr("key", []byte("next_prefix")).AddInt64("val", int64(nextPrefix+uint32(len(tdef.Indexes))))

    _, err = tx.Upsert("@meta", *metaKey)
    return err
}