fix(dbml): honor composite [pk] in Indexes blocks, preserve column order

A composite [pk] entry inside an Indexes block (e.g. (a, b) [pk]) was silently dropped: models.Index has no way to represent a primary key, so the attribute was parsed and ignored, producing neither a PK nor a meaningful index. It's now converted into a PrimaryKeyConstraint. Also, Column.Sequence was never set by the DBML reader, so composite PKs assembled from column-level [pk] attributes fell back to alphabetical Name sorting instead of declaration order. Columns now get a per-table sequence counter reflecting the order they were declared.
2026-08-24 12:59:18 +02:00
parent 241bfc2302
commit 7fb343596a
2 changed files with 172 additions and 23 deletions
@@ -434,6 +434,7 @@ func (r *Reader) parseDBML(content string) (*models.Database, error) {
 	var currentSchema string
 	var inIndexes bool
 	var inTable bool
+	var columnSeq uint

 	tableRegex := regexp.MustCompile(`^Table\s+(.+?)\s*{`)
 	refRegex := regexp.MustCompile(`^Ref:\s+(.+)`)
@@ -469,6 +470,7 @@ func (r *Reader) parseDBML(content string) (*models.Database, error) {
 			currentTable = models.InitTable(tableName, currentSchema)
 			inTable = true
 			inIndexes = false
+			columnSeq = 0
 			continue
 		}

@@ -497,6 +499,17 @@ func (r *Reader) parseDBML(content string) (*models.Database, error) {

 		// Parse index definition
 		if inIndexes && currentTable != nil {
+			// A composite `[pk]` entry inside an Indexes block declares the
+			// table's primary key (DBML's way of expressing multi-column PKs
+			// that can't be attached to a single column). It must become a
+			// primary key constraint, not a plain index, or the PK is lost.
+			if indexLineHasPKAttr(line) {
+				if constraint := r.parsePrimaryKeyIndex(line, currentTable.Name, currentSchema); constraint != nil {
+					currentTable.Constraints[constraint.Name] = constraint
+				}
+				continue
+			}
+
 			index := r.parseIndex(line, currentTable.Name, currentSchema)
 			if index != nil {
 				currentTable.Indexes[index.Name] = index
@@ -516,6 +529,8 @@ func (r *Reader) parseDBML(content string) (*models.Database, error) {
 		if inTable && !inIndexes && currentTable != nil {
 			column, constraint := r.parseColumn(line, currentTable.Name, currentSchema)
 			if column != nil {
+				columnSeq++
+				column.Sequence = columnSeq
 				currentTable.Columns[column.Name] = column
 			}
 			if constraint != nil {
@@ -743,9 +758,10 @@ func stripWrappingQuotes(s string) string {
 	return s
 }

-// parseIndex parses a DBML index definition
-func (r *Reader) parseIndex(line, tableName, schemaName string) *models.Index {
-	// Format: (columns) [attributes] OR columnname [attributes]
+// indexLineColumns extracts the column list from an Indexes-block entry,
+// e.g. "(col1, col2) [attrs]" or "columnname [attrs]", preserving
+// declaration order.
+func indexLineColumns(line string) []string {
 	var columns []string

 	// Find the attributes section to avoid parsing parentheses in notes/attributes
@@ -776,6 +792,56 @@ func (r *Reader) parseIndex(line, tableName, schemaName string) *models.Index {
 		}
 	}

+	return columns
+}
+
+// indexLineAttrs extracts and splits the bracketed attribute list of an
+// Indexes-block entry, e.g. "[pk]" or "[unique, name: 'foo']".
+func indexLineAttrs(line string) []string {
+	attrStart := strings.Index(line, "[")
+	attrEnd := strings.Index(line, "]")
+	if attrStart < 0 || attrEnd < 0 || attrStart >= attrEnd {
+		return nil
+	}
+
+	var attrs []string
+	for _, attr := range strings.Split(line[attrStart+1:attrEnd], ",") {
+		attrs = append(attrs, strings.TrimSpace(attr))
+	}
+	return attrs
+}
+
+// indexLineHasPKAttr reports whether an Indexes-block entry carries a `pk`
+// attribute, e.g. "(artifact_id, sha256) [pk]". DBML uses this form to
+// declare composite primary keys that can't be attached to a single column.
+func indexLineHasPKAttr(line string) bool {
+	for _, attr := range indexLineAttrs(line) {
+		if attr == "pk" || attr == "primary key" {
+			return true
+		}
+	}
+	return false
+}
+
+// parsePrimaryKeyIndex converts a composite `[pk]` entry from an Indexes
+// block into a primary key constraint, preserving the declared column order.
+func (r *Reader) parsePrimaryKeyIndex(line, tableName, schemaName string) *models.Constraint {
+	columns := indexLineColumns(line)
+	if len(columns) == 0 {
+		return nil
+	}
+
+	constraint := models.InitConstraint("pk_"+tableName, models.PrimaryKeyConstraint)
+	constraint.Schema = schemaName
+	constraint.Table = tableName
+	constraint.Columns = columns
+	return constraint
+}
+
+// parseIndex parses a DBML index definition
+func (r *Reader) parseIndex(line, tableName, schemaName string) *models.Index {
+	// Format: (columns) [attributes] OR columnname [attributes]
+	columns := indexLineColumns(line)
 	if len(columns) == 0 {
 		return nil
 	}
@@ -786,26 +852,15 @@ func (r *Reader) parseIndex(line, tableName, schemaName string) *models.Index {
 	index.Columns = columns

 	// Parse attributes
-	if strings.Contains(line, "[") && strings.Contains(line, "]") {
-		attrStart := strings.Index(line, "[")
-		attrEnd := strings.Index(line, "]")
-		if attrStart < attrEnd {
-			attrs := line[attrStart+1 : attrEnd]
-			attrList := strings.Split(attrs, ",")
-
-			for _, attr := range attrList {
-				attr = strings.TrimSpace(attr)
-
-				if attr == "unique" {
-					index.Unique = true
-				} else if strings.HasPrefix(attr, "name:") {
-					name := strings.TrimSpace(strings.TrimPrefix(attr, "name:"))
-					index.Name = strings.Trim(name, "'\"")
-				} else if strings.HasPrefix(attr, "type:") {
-					indexType := strings.TrimSpace(strings.TrimPrefix(attr, "type:"))
-					index.Type = strings.Trim(indexType, "'\"")
-				}
-			}
+	for _, attr := range indexLineAttrs(line) {
+		if attr == "unique" {
+			index.Unique = true
+		} else if strings.HasPrefix(attr, "name:") {
+			name := strings.TrimSpace(strings.TrimPrefix(attr, "name:"))
+			index.Name = strings.Trim(name, "'\"")
+		} else if strings.HasPrefix(attr, "type:") {
+			indexType := strings.TrimSpace(strings.TrimPrefix(attr, "type:"))
+			index.Type = strings.Trim(indexType, "'\"")
 		}
 	}

@@ -932,3 +932,97 @@ func TestHasCommentedRefs(t *testing.T) {
 		})
 	}
 }
+
+// TestReader_CompositePKIndex verifies that a composite `[pk]` entry inside
+// an Indexes block is turned into a primary key constraint, in declaration
+// order, rather than being silently dropped.
+func TestReader_CompositePKIndex(t *testing.T) {
+	dbmlContent := `Table artifact_blob {
+  artifact_id integer [not null]
+  sha256 text [not null]
+  size integer
+
+  Indexes {
+    (artifact_id, sha256) [pk]
+  }
+}
+`
+	dir := t.TempDir()
+	path := filepath.Join(dir, "composite_pk.dbml")
+	if err := os.WriteFile(path, []byte(dbmlContent), 0644); err != nil {
+		t.Fatalf("failed to write fixture: %v", err)
+	}
+
+	reader := NewReader(&readers.ReaderOptions{FilePath: path})
+	db, err := reader.ReadDatabase()
+	if err != nil {
+		t.Fatalf("ReadDatabase() error = %v", err)
+	}
+
+	table := db.Schemas[0].Tables[0]
+
+	var pk *models.Constraint
+	for _, c := range table.Constraints {
+		if c.Type == models.PrimaryKeyConstraint {
+			pk = c
+			break
+		}
+	}
+	if pk == nil {
+		t.Fatal("expected a primary key constraint, got none")
+	}
+	want := []string{"artifact_id", "sha256"}
+	if len(pk.Columns) != len(want) {
+		t.Fatalf("expected PK columns %v, got %v", want, pk.Columns)
+	}
+	for i, col := range want {
+		if pk.Columns[i] != col {
+			t.Errorf("PK column[%d] = %q, want %q (order must match declaration)", i, pk.Columns[i], col)
+		}
+	}
+
+	// No plain index should be emitted for the pk-only entry.
+	if len(table.Indexes) != 0 {
+		t.Errorf("expected no plain indexes from a [pk] Indexes entry, got %v", table.Indexes)
+	}
+}
+
+// TestReader_ColumnPKOrderPreserved verifies that composite primary keys
+// declared via column-level [pk] attributes keep declaration order (via
+// Column.Sequence) instead of falling back to alphabetical sorting.
+func TestReader_ColumnPKOrderPreserved(t *testing.T) {
+	dbmlContent := `Table snapshot_artifact {
+  snapshot_id integer [pk, not null]
+  artifact_id integer [pk, not null]
+}
+`
+	dir := t.TempDir()
+	path := filepath.Join(dir, "column_pk_order.dbml")
+	if err := os.WriteFile(path, []byte(dbmlContent), 0644); err != nil {
+		t.Fatalf("failed to write fixture: %v", err)
+	}
+
+	reader := NewReader(&readers.ReaderOptions{FilePath: path})
+	db, err := reader.ReadDatabase()
+	if err != nil {
+		t.Fatalf("ReadDatabase() error = %v", err)
+	}
+
+	table := db.Schemas[0].Tables[0]
+
+	snapshotCol, ok := table.Columns["snapshot_id"]
+	if !ok {
+		t.Fatal("column 'snapshot_id' not found")
+	}
+	artifactCol, ok := table.Columns["artifact_id"]
+	if !ok {
+		t.Fatal("column 'artifact_id' not found")
+	}
+
+	if snapshotCol.Sequence == 0 || artifactCol.Sequence == 0 {
+		t.Fatalf("expected non-zero Sequence values, got snapshot_id=%d artifact_id=%d", snapshotCol.Sequence, artifactCol.Sequence)
+	}
+	if snapshotCol.Sequence >= artifactCol.Sequence {
+		t.Errorf("expected snapshot_id (declared first) to have a lower Sequence than artifact_id, got %d >= %d", snapshotCol.Sequence, artifactCol.Sequence)
+	}
+}