SQLite: Adding support for composite keys.

sqlite/_dumps/structs.sql

@@ -62,4 +62,13 @@ CREATE TABLE stats_test (
   value integer
 );
 
+DROP TABLE IF EXISTS composite_keys;
+
+CREATE TABLE composite_keys (
+  code VARCHAR(255) default '',
+  user_id VARCHAR(255) default '',
+  some_val VARCHAR(255) default '',
+  primary key (code, user_id)
+);
+
 COMMIT;

sqlite/collection.go

@@ -82,6 +82,14 @@ func whereValues(term interface{}) (where sqlgen.Where, args []interface{}) {
 		for _, kk := range k {
 			where = append(where, kk)
 		}
+	case db.Constrainer:
+		k, v := conditionValues(t.Constraint())
+		args = append(args, v...)
+		for _, kk := range k {
+			where = append(where, kk)
+		}
+	default:
+		panic(fmt.Sprintf(db.ErrUnknownConditionType.Error(), reflect.TypeOf(t)))
 	}
 
 	return where, args
@@ -215,9 +223,12 @@ func (c *table) Truncate() error {
 
 // Appends an item (map or struct) into the collection.
 func (c *table) Append(item interface{}) (interface{}, error) {
-	var arguments []interface{}
+
+	var pKey []string
 	var columns sqlgen.Columns
 	var values sqlgen.Values
+	var arguments []interface{}
+	var id []interface{}
 
 	cols, vals, err := c.FieldValues(item, toInternal)
 
@@ -245,19 +256,55 @@ func (c *table) Append(item interface{}) (interface{}, error) {
 		}
 	}
 
-	row, err := c.source.doExec(sqlgen.Statement{
+	if pKey, err = c.source.getPrimaryKey(c.tableN(0)); err != nil {
+		if err != sql.ErrNoRows {
+			// Can't tell primary key.
+			return nil, err
+		}
+	}
+
+	stmt := sqlgen.Statement{
 		Type:    sqlgen.SqlInsert,
 		Table:   sqlgen.Table{c.tableN(0)},
 		Columns: columns,
 		Values:  values,
-	}, arguments...)
+	}
 
-	if err != nil {
+	var res sql.Result
+	if res, err = c.source.doExec(stmt, arguments...); err != nil {
 		return nil, err
 	}
 
-	var id int64
-	id, _ = row.LastInsertId()
+	if len(pKey) == 1 {
+		lastID, _ := res.LastInsertId()
+		id = []interface{}{lastID}
+	} else {
+		// There is no "RETURNING" in SQLite, so we have to return the values that
+		// were given for constructing the composite key.
+		id = make([]interface{}, len(pKey))
+
+		for i := range cols {
+			for j := 0; j < len(pKey); j++ {
+				if pKey[j] == cols[i] {
+					id[j] = vals[i]
+				}
+			}
+		}
+	}
+
+	// Does the item satisfy the db.ID interface?
+	if setter, ok := item.(db.IDSetter); ok {
+		if err := setter.SetID(id...); err != nil {
+			return nil, err
+		}
+	}
+
+	// Backwards compatibility (int64).
+	if len(id) == 1 {
+		if numericID, ok := id[0].(int64); ok {
+			return numericID, nil
+		}
+	}
 
 	return id, nil
 }

sqlite/database.go

@@ -59,10 +59,15 @@ type source struct {
 	session *sql.DB
 	tx      *tx
 	schema  *schema.DatabaseSchema
+	// columns property was introduced so we could query PRAGMA data only once
+	// and retrieve all the column information we'd need, such as name and if it
+	// is a primary key.
+	columns []columnSchemaT
 }
 
 type columnSchemaT struct {
 	Name string `db:"name"`
+	PK   int    `db:"pk"`
 }
 
 func debugEnabled() bool {
@@ -477,16 +482,16 @@ func (s *source) tableColumns(tableName string) ([]string, error) {
 
 	rows, err := s.doRawQuery(q)
 
-	tableFields := []columnSchemaT{}
+	s.columns = []columnSchemaT{}
 
-	if err = sqlutil.FetchRows(rows, &tableFields); err != nil {
+	if err = sqlutil.FetchRows(rows, &s.columns); err != nil {
 		return nil, err
 	}
 
-	s.schema.TableInfo[tableName].Columns = make([]string, 0, len(tableFields))
+	s.schema.TableInfo[tableName].Columns = make([]string, 0, len(s.columns))
 
-	for i := range tableFields {
-		s.schema.TableInfo[tableName].Columns = append(s.schema.TableInfo[tableName].Columns, tableFields[i].Name)
+	for i := range s.columns {
+		s.schema.TableInfo[tableName].Columns = append(s.schema.TableInfo[tableName].Columns, s.columns[i].Name)
 	}
 
 	return s.schema.TableInfo[tableName].Columns, nil
@@ -531,3 +536,29 @@ func (s *source) Collection(names ...string) (db.Collection, error) {
 
 	return col, nil
 }
+
+// getPrimaryKey returns the names of the columns that define the primary key
+// of the table.
+func (s *source) getPrimaryKey(tableName string) ([]string, error) {
+	tableSchema := s.schema.Table(tableName)
+
+	s.tableColumns(tableName)
+
+	maxValue := -1
+
+	for i := range s.columns {
+		if s.columns[i].PK > 0 && s.columns[i].PK > maxValue {
+			maxValue = s.columns[i].PK
+		}
+	}
+
+	tableSchema.PrimaryKey = make([]string, maxValue)
+
+	for i := range s.columns {
+		if s.columns[i].PK > 0 {
+			tableSchema.PrimaryKey[s.columns[i].PK-1] = s.columns[i].Name
+		}
+	}
+
+	return tableSchema.PrimaryKey, nil
+}

sqlite/database_test.go

@@ -30,9 +30,11 @@ package sqlite
 
 import (
 	"database/sql"
+	"errors"
 	"math/rand"
 	"os"
 	"reflect"
+	"strconv"
 	"strings"
 	"testing"
 	"time"
@@ -76,6 +78,29 @@ type testValuesStruct struct {
 	Time  time.Duration `field:"_time"`
 }
 
+type ItemWithKey struct {
+	Code    string `db:"code"`
+	UserID  string `db:"user_id"`
+	SomeVal string `db:"some_val"`
+}
+
+func (item ItemWithKey) Constraint() db.Cond {
+	cond := db.Cond{
+		"code":    item.Code,
+		"user_id": item.UserID,
+	}
+	return cond
+}
+
+func (item *ItemWithKey) SetID(keys ...interface{}) error {
+	if len(keys) == 2 {
+		item.Code = keys[0].(string)
+		item.UserID = keys[1].(string)
+		return nil
+	}
+	return errors.New(`Expecting exactly two keys.`)
+}
+
 var testValues testValuesStruct
 
 func init() {
@@ -1208,6 +1233,64 @@ func TestTransactionsAndRollback(t *testing.T) {
 
 }
 
+// Attempts to test composite keys.
+func TestCompositeKeys(t *testing.T) {
+	var err error
+	var id interface{}
+	var sess db.Database
+	var compositeKeys db.Collection
+
+	if sess, err = db.Open(Adapter, settings); err != nil {
+		t.Fatal(err)
+	}
+
+	defer sess.Close()
+
+	if compositeKeys, err = sess.Collection("composite_keys"); err != nil {
+		t.Fatal(err)
+	}
+
+	n := rand.Intn(100000)
+
+	item := ItemWithKey{
+		"ABCDEF",
+		strconv.Itoa(n),
+		"Some value",
+	}
+
+	if id, err = compositeKeys.Append(&item); err != nil {
+		t.Fatal(err)
+	}
+
+	ids := id.([]interface{})
+
+	if ids[0].(string) != item.Code {
+		t.Fatal(`Keys must match.`)
+	}
+
+	if ids[1].(string) != item.UserID {
+		t.Fatal(`Keys must match.`)
+	}
+
+	// Using constraint interface.
+	res := compositeKeys.Find(ItemWithKey{Code: item.Code, UserID: item.UserID})
+
+	var item2 ItemWithKey
+
+	if item2.SomeVal == item.SomeVal {
+		t.Fatal(`Values must be different before query.`)
+	}
+
+	if err := res.One(&item2); err != nil {
+		t.Fatal(err)
+	}
+
+	if item2.SomeVal != item.SomeVal {
+		t.Fatal(`Values must be equal after query.`)
+	}
+
+}
+
 // Attempts to add many different datatypes to a single row in a collection,
 // then it tries to get the stored datatypes and check if the stored and the
 // original values match.