respect primary keys ordering (cloudspannerecosystem#44)

internal/loader.go

@@ -130,12 +130,47 @@ func (tl *TypeLoader) LoadTable(args *ArgType) (map[string]*Type, error) {
 			return nil, err
 		}
 
+		if err := tl.loadPrimaryKeys(typeTpl); err != nil {
+			return nil, err
+		}
+
 		tableMap[ti.TableName] = typeTpl
 	}
 
 	return tableMap, nil
 }
 
+// loadPrimaryKeys loads primary key fields
+func (tl *TypeLoader) loadPrimaryKeys(typeTpl *Type) error {
+	// reorder primary keys
+	indexCols, err := tl.loader.IndexColumnList(typeTpl.Table.TableName, "PRIMARY_KEY")
+	if err != nil {
+		panic(err)
+	}
+
+	var fields []*Field
+	for _, idx := range indexCols {
+		var field *Field
+		for _, f := range typeTpl.Fields {
+			if f.Col.ColumnName == idx.ColumnName {
+				field = f
+				break
+			}
+		}
+
+		if field == nil {
+			return fmt.Errorf("primary key column is not found in column list: table=%v column=%v",
+				typeTpl.Name, idx.ColumnName,
+			)
+		}
+		fields = append(fields, field)
+	}
+
+	typeTpl.PrimaryKey = fields[0] // backward compatibility
+	typeTpl.PrimaryKeyFields = fields
+	return nil
+}
+
 // tableCustomTypes find custom type definitions of the table
 func (tl *TypeLoader) tableCustomTypes(table string) map[string]string {
 	var columnTypes map[string]string
@@ -198,13 +233,6 @@ func (tl *TypeLoader) LoadColumns(args *ArgType, typeTpl *Type) error {
 			}
 		}
 
-		// set primary key
-		if c.IsPrimaryKey {
-			typeTpl.PrimaryKeyFields = append(typeTpl.PrimaryKeyFields, f)
-			// This is retained for backward compatibility in the templates.
-			typeTpl.PrimaryKey = f
-		}
-
 		// append col to template fields
 		typeTpl.Fields = append(typeTpl.Fields, f)
 	}
@@ -264,11 +292,6 @@ func (tl *TypeLoader) LoadTableIndexes(args *ArgType, typeTpl *Type, ixMap map[s
 		ixMap[typeTpl.Table.TableName+"_"+ix.IndexName] = ixTpl
 	}
 
-	// search for primary key if it was skipped being set in the type
-	if typeTpl.PrimaryKey == nil {
-		return fmt.Errorf("no primary key found for %v", typeTpl.Name)
-	}
-
 	return nil
 }
 

loaders/parser.go

@@ -142,6 +142,10 @@ func (s *SpannerLoaderFromDDL) IndexList(name string) ([]*models.Index, error) {
 }
 
 func (s *SpannerLoaderFromDDL) IndexColumnList(table, index string) ([]*models.IndexColumn, error) {
+	if index == "PRIMARY_KEY" {
+		return s.primaryKeyColumnList(table)
+	}
+
 	var cols []*models.IndexColumn
 	for _, ix := range s.tables[table].createIndexes {
 		if ix.Name.Name != index {
@@ -170,3 +174,20 @@ func (s *SpannerLoaderFromDDL) IndexColumnList(table, index string) ([]*models.I
 
 	return cols, nil
 }
+
+func (s *SpannerLoaderFromDDL) primaryKeyColumnList(table string) ([]*models.IndexColumn, error) {
+	tbl, ok := s.tables[table]
+	if !ok {
+		return nil, nil
+	}
+
+	var cols []*models.IndexColumn
+	for i, key := range tbl.createTable.PrimaryKeys {
+		cols = append(cols, &models.IndexColumn{
+			SeqNo:      i + 1,
+			ColumnName: key.Name.Name,
+		})
+	}
+
+	return cols, nil
+}

test/testdata/schema.sql

@@ -32,6 +32,12 @@ CREATE INDEX CompositePrimaryKeysByError  ON CompositePrimaryKeys(Error);
 CREATE INDEX CompositePrimaryKeysByError2 ON CompositePrimaryKeys(Error) STORING(Z);
 CREATE INDEX CompositePrimaryKeysByError3 ON CompositePrimaryKeys(Error) STORING(Z, Y);
 
+CREATE TABLE OutOfOrderPrimaryKeys (
+  PKey1 STRING(32) NOT NULL,
+  PKey2 STRING(32) NOT NULL,
+  PKey3 STRING(32) NOT NULL,
+) PRIMARY KEY(PKey2, PKey1, PKey3);
+
 CREATE TABLE FullTypes (
   PKey STRING(32) NOT NULL,
   FTString STRING(32) NOT NULL,