From 7fb343596a0506944a968b4d41e17402dae24c57 Mon Sep 17 00:00:00 2001 From: Hein Date: Mon, 24 Aug 2026 12:59:18 +0200 Subject: [PATCH] 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. --- pkg/readers/dbml/reader.go | 101 ++++++++++++++++++++++++-------- pkg/readers/dbml/reader_test.go | 94 +++++++++++++++++++++++++++++ 2 files changed, 172 insertions(+), 23 deletions(-) diff --git a/pkg/readers/dbml/reader.go b/pkg/readers/dbml/reader.go index c89a752..f9cfe08 100644 --- a/pkg/readers/dbml/reader.go +++ b/pkg/readers/dbml/reader.go @@ -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, "'\"") } } diff --git a/pkg/readers/dbml/reader_test.go b/pkg/readers/dbml/reader_test.go index e9e39e3..463bdaa 100644 --- a/pkg/readers/dbml/reader_test.go +++ b/pkg/readers/dbml/reader_test.go @@ -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) + } +}