sql writer

pkg/readers/dbml/reader.go

@@ -67,6 +67,15 @@ func (r *Reader) ReadTable() (*models.Table, error) {
 	return schema.Tables[0], nil
 }
 
+// stripQuotes removes surrounding quotes from an identifier
+func stripQuotes(s string) string {
+	s = strings.TrimSpace(s)
+	if len(s) >= 2 && ((s[0] == '"' && s[len(s)-1] == '"') || (s[0] == '\'' && s[len(s)-1] == '\'')) {
+		return s[1 : len(s)-1]
+	}
+	return s
+}
+
 // parseDBML parses DBML content and returns a Database model
 func (r *Reader) parseDBML(content string) (*models.Database, error) {
 	db := models.InitDatabase("database")
@@ -79,13 +88,14 @@ func (r *Reader) parseDBML(content string) (*models.Database, error) {
 
 	scanner := bufio.NewScanner(strings.NewReader(content))
 	schemaMap := make(map[string]*models.Schema)
+	pendingConstraints := []*models.Constraint{}
 
 	var currentTable *models.Table
 	var currentSchema string
 	var inIndexes bool
 	var inTable bool
 
-	tableRegex := regexp.MustCompile(`^Table\s+([a-zA-Z0-9_.]+)\s*{`)
+	tableRegex := regexp.MustCompile(`^Table\s+(.+?)\s*{`)
 	refRegex := regexp.MustCompile(`^Ref:\s+(.+)`)
 
 	for scanner.Scan() {
@@ -102,10 +112,11 @@ func (r *Reader) parseDBML(content string) (*models.Database, error) {
 			parts := strings.Split(tableName, ".")
 
 			if len(parts) == 2 {
-				currentSchema = parts[0]
-				tableName = parts[1]
+				currentSchema = stripQuotes(parts[0])
+				tableName = stripQuotes(parts[1])
 			} else {
 				currentSchema = "public"
+				tableName = stripQuotes(parts[0])
 			}
 
 			// Ensure schema exists
@@ -131,7 +142,7 @@ func (r *Reader) parseDBML(content string) (*models.Database, error) {
 		}
 
 		// Parse indexes section
-		if inTable && strings.HasPrefix(line, "indexes") {
+		if inTable && (strings.HasPrefix(line, "Indexes {") || strings.HasPrefix(line, "indexes {")) {
 			inIndexes = true
 			continue
 		}
@@ -161,10 +172,14 @@ func (r *Reader) parseDBML(content string) (*models.Database, error) {
 
 		// Parse column definition
 		if inTable && !inIndexes && currentTable != nil {
-			column := r.parseColumn(line, currentTable.Name, currentSchema)
+			column, constraint := r.parseColumn(line, currentTable.Name, currentSchema)
 			if column != nil {
 				currentTable.Columns[column.Name] = column
 			}
+			if constraint != nil {
+				// Add to pending list - will assign to tables at the end
+				pendingConstraints = append(pendingConstraints, constraint)
+			}
 			continue
 		}
 
@@ -186,6 +201,19 @@ func (r *Reader) parseDBML(content string) (*models.Database, error) {
 		}
 	}
 
+	// Assign pending constraints to their respective tables
+	for _, constraint := range pendingConstraints {
+		// Find the table this constraint belongs to
+		if schema, exists := schemaMap[constraint.Schema]; exists {
+			for _, table := range schema.Tables {
+				if table.Name == constraint.Table {
+					table.Constraints[constraint.Name] = constraint
+					break
+				}
+			}
+		}
+	}
+
 	// Add schemas to database
 	for _, schema := range schemaMap {
 		db.Schemas = append(db.Schemas, schema)
@@ -195,19 +223,21 @@ func (r *Reader) parseDBML(content string) (*models.Database, error) {
 }
 
 // parseColumn parses a DBML column definition
-func (r *Reader) parseColumn(line, tableName, schemaName string) *models.Column {
+func (r *Reader) parseColumn(line, tableName, schemaName string) (*models.Column, *models.Constraint) {
 	// Format: column_name type [attributes] // comment
 	parts := strings.Fields(line)
 	if len(parts) < 2 {
-		return nil
+		return nil, nil
 	}
 
-	columnName := parts[0]
-	columnType := parts[1]
+	columnName := stripQuotes(parts[0])
+	columnType := stripQuotes(parts[1])
 
 	column := models.InitColumn(columnName, tableName, schemaName)
 	column.Type = columnType
 
+	var constraint *models.Constraint
+
 	// Parse attributes in brackets
 	if strings.Contains(line, "[") && strings.Contains(line, "]") {
 		attrStart := strings.Index(line, "[")
@@ -230,7 +260,55 @@ func (r *Reader) parseColumn(line, tableName, schemaName string) *models.Column
 					defaultVal := strings.TrimSpace(strings.TrimPrefix(attr, "default:"))
 					column.Default = strings.Trim(defaultVal, "'\"")
 				} else if attr == "unique" {
-					// Could create a unique constraint here
+					// Create a unique constraint
+					uniqueConstraint := models.InitConstraint(
+						fmt.Sprintf("uq_%s", columnName),
+						models.UniqueConstraint,
+					)
+					uniqueConstraint.Schema = schemaName
+					uniqueConstraint.Table = tableName
+					uniqueConstraint.Columns = []string{columnName}
+					// Store it to be added later
+					if constraint == nil {
+						constraint = uniqueConstraint
+					}
+				} else if strings.HasPrefix(attr, "ref:") {
+					// Parse inline reference
+					// DBML semantics depend on context:
+					// - On FK column: ref: < target means "this FK references target"
+					// - On PK column: ref: < source means "source references this PK" (reverse notation)
+					refStr := strings.TrimSpace(strings.TrimPrefix(attr, "ref:"))
+
+					// Check relationship direction operator
+					refOp := strings.TrimSpace(refStr)
+					var isReverse bool
+					if strings.HasPrefix(refOp, "<") {
+						isReverse = column.IsPrimaryKey // < on PK means "is referenced by" (reverse)
+					} else if strings.HasPrefix(refOp, ">") {
+						isReverse = !column.IsPrimaryKey // > on FK means reverse
+					}
+
+					constraint = r.parseRef(refStr)
+					if constraint != nil {
+						if isReverse {
+							// Reverse: parsed ref is SOURCE, current column is TARGET
+							// Constraint should be ON the source table
+							constraint.Schema = constraint.ReferencedSchema
+							constraint.Table = constraint.ReferencedTable
+							constraint.Columns = constraint.ReferencedColumns
+							constraint.ReferencedSchema = schemaName
+							constraint.ReferencedTable = tableName
+							constraint.ReferencedColumns = []string{columnName}
+						} else {
+							// Forward: current column is SOURCE, parsed ref is TARGET
+							// Standard FK: constraint is ON current table
+							constraint.Schema = schemaName
+							constraint.Table = tableName
+							constraint.Columns = []string{columnName}
+						}
+						// Generate short constraint name based on the column
+						constraint.Name = fmt.Sprintf("fk_%s", constraint.Columns[0])
+					}
 				}
 			}
 		}
@@ -242,28 +320,41 @@ func (r *Reader) parseColumn(line, tableName, schemaName string) *models.Column
 		column.Comment = strings.TrimSpace(line[commentStart+2:])
 	}
 
-	return column
+	return column, constraint
 }
 
 // parseIndex parses a DBML index definition
 func (r *Reader) parseIndex(line, tableName, schemaName string) *models.Index {
-	// Format: (columns) [attributes]
-	if !strings.Contains(line, "(") || !strings.Contains(line, ")") {
-		return nil
+	// Format: (columns) [attributes] OR columnname [attributes]
+	var columns []string
+
+	if strings.Contains(line, "(") && strings.Contains(line, ")") {
+		// Multi-column format: (col1, col2) [attributes]
+		colStart := strings.Index(line, "(")
+		colEnd := strings.Index(line, ")")
+		if colStart >= colEnd {
+			return nil
+		}
+
+		columnsStr := line[colStart+1 : colEnd]
+		for _, col := range strings.Split(columnsStr, ",") {
+			columns = append(columns, stripQuotes(strings.TrimSpace(col)))
+		}
+	} else {
+		// Single column format: columnname [attributes]
+		// Extract column name before the bracket
+		if strings.Contains(line, "[") {
+			colName := strings.TrimSpace(line[:strings.Index(line, "[")])
+			if colName != "" {
+				columns = []string{stripQuotes(colName)}
+			}
+		}
 	}
 
-	colStart := strings.Index(line, "(")
-	colEnd := strings.Index(line, ")")
-	if colStart >= colEnd {
+	if len(columns) == 0 {
 		return nil
 	}
 
-	columnsStr := line[colStart+1 : colEnd]
-	columns := strings.Split(columnsStr, ",")
-	for i := range columns {
-		columns[i] = strings.TrimSpace(columns[i])
-	}
-
 	index := models.InitIndex("")
 	index.Table = tableName
 	index.Schema = schemaName
@@ -304,9 +395,11 @@ func (r *Reader) parseIndex(line, tableName, schemaName string) *models.Index {
 // parseRef parses a DBML Ref (foreign key relationship)
 func (r *Reader) parseRef(refStr string) *models.Constraint {
 	// Format: schema.table.(columns) > schema.table.(columns) [actions]
+	// Or inline format: < schema.table.column (for inline column refs)
 
 	// Split by relationship operator (>, <, -, etc.)
 	var fromPart, toPart string
+	isInlineRef := false
 
 	for _, op := range []string{">", "<", "-"} {
 		if strings.Contains(refStr, op) {
@@ -314,30 +407,53 @@ func (r *Reader) parseRef(refStr string) *models.Constraint {
 			if len(parts) == 2 {
 				fromPart = strings.TrimSpace(parts[0])
 				toPart = strings.TrimSpace(parts[1])
+				// Check if this is an inline ref (operator at start)
+				if fromPart == "" {
+					isInlineRef = true
+				}
 				break
 			}
 		}
 	}
 
-	if fromPart == "" || toPart == "" {
+	// For inline refs, only toPart should be populated
+	if isInlineRef {
+		if toPart == "" {
+			return nil
+		}
+	} else if fromPart == "" || toPart == "" {
 		return nil
 	}
 
 	// Remove actions part if present
-	if strings.Contains(toPart, "[") {
-		toPart = strings.TrimSpace(toPart[:strings.Index(toPart, "[")])
+	if idx := strings.Index(toPart, "["); idx >= 0 {
+		toPart = strings.TrimSpace(toPart[:idx])
 	}
 
-	// Parse from table and column
-	fromSchema, fromTable, fromColumns := r.parseTableRef(fromPart)
+	// Parse references
+	var fromSchema, fromTable string
+	var fromColumns []string
 	toSchema, toTable, toColumns := r.parseTableRef(toPart)
 
-	if fromTable == "" || toTable == "" {
+	if !isInlineRef {
+		fromSchema, fromTable, fromColumns = r.parseTableRef(fromPart)
+		if fromTable == "" {
+			return nil
+		}
+	}
+
+	if toTable == "" {
 		return nil
 	}
 
+	// Generate short constraint name based on the source column
+	constraintName := fmt.Sprintf("fk_%s_%s", fromTable, toTable)
+	if len(fromColumns) > 0 {
+		constraintName = fmt.Sprintf("fk_%s", fromColumns[0])
+	}
+
 	constraint := models.InitConstraint(
-		fmt.Sprintf("fk_%s_%s", fromTable, toTable),
+		constraintName,
 		models.ForeignKeyConstraint,
 	)
 
@@ -371,29 +487,48 @@ func (r *Reader) parseRef(refStr string) *models.Constraint {
 	return constraint
 }
 
-// parseTableRef parses a table reference like "schema.table.(column1, column2)"
+// parseTableRef parses a table reference like "schema.table.(column1, column2)" or "schema"."table"."column"
 func (r *Reader) parseTableRef(ref string) (schema, table string, columns []string) {
-	// Extract columns if present
+	// Extract columns if present in parentheses format
+	hasParentheses := false
 	if strings.Contains(ref, "(") && strings.Contains(ref, ")") {
 		colStart := strings.Index(ref, "(")
 		colEnd := strings.Index(ref, ")")
 		if colStart < colEnd {
 			columnsStr := ref[colStart+1 : colEnd]
 			for _, col := range strings.Split(columnsStr, ",") {
-				columns = append(columns, strings.TrimSpace(col))
+				columns = append(columns, stripQuotes(strings.TrimSpace(col)))
 			}
+			hasParentheses = true
 		}
 		ref = ref[:colStart]
 	}
 
-	// Parse schema and table
+	// Parse schema, table, and optionally column
 	parts := strings.Split(strings.TrimSpace(ref), ".")
-	if len(parts) == 2 {
-		schema = parts[0]
-		table = parts[1]
+	if len(parts) == 3 {
+		// Format: "schema"."table"."column"
+		schema = stripQuotes(parts[0])
+		table = stripQuotes(parts[1])
+		if !hasParentheses {
+			columns = []string{stripQuotes(parts[2])}
+		}
+	} else if len(parts) == 2 {
+		// Could be "schema"."table" or "table"."column"
+		// If columns are already extracted from parentheses, this is schema.table
+		// If no parentheses, this is table.column
+		if hasParentheses {
+			schema = stripQuotes(parts[0])
+			table = stripQuotes(parts[1])
+		} else {
+			schema = "public"
+			table = stripQuotes(parts[0])
+			columns = []string{stripQuotes(parts[1])}
+		}
 	} else if len(parts) == 1 {
+		// Format: "table"
 		schema = "public"
-		table = parts[0]
+		table = stripQuotes(parts[0])
 	}
 
 	return