Allow column constraints to reference themselves

Author: wickwirew

Sources/Compiler/Compiler.swift

@@ -7,46 +7,32 @@
 
 public struct Compiler {
     public var schema = Schema()
-    public private(set) var queries: [Statement] = []
-    public private(set) var migrations: [Statement] = []
-    public private(set) var diagnostics = Diagnostics()
-    
     private var pragmas = PragmaAnalyzer()
 
     public init() {}
 
-    public var hasDiagnostics: Bool {
-        return !diagnostics.isEmpty
-    }
-    
-    public mutating func compile(migration: String) -> Diagnostics {
-        let (stmts, diagnostics) = compile(
+    public mutating func compile(migration: String) -> ([Statement], Diagnostics) {
+        compile(
             source: migration,
             validator: IsValidForMigrations(),
             context: "migrations"
         )
-        self.migrations.append(contentsOf: stmts)
-        self.diagnostics.merge(diagnostics)
-        return diagnostics
     }
 
-    public mutating func compile(queries: String) -> Diagnostics {
-        let (stmts, diagnostics) = compile(
+    public mutating func compile(queries: String) -> ([Statement], Diagnostics) {
+        compile(
             source: queries,
             validator: IsValidForQueries(),
             context: "queries"
         )
-        self.queries.append(contentsOf: stmts)
-        self.diagnostics.merge(diagnostics)
-        return diagnostics
     }
 
     public mutating func compile(
         query: String,
         named name: String,
         inputType: String?,
         outputType: String?
-    ) -> Diagnostics {
+    ) -> (Statement?, Diagnostics) {
         var (stmts, diagnostics) = compile(
             source: query,
             validator: IsValidForQueries(),
@@ -56,8 +42,7 @@ public struct Compiler {
         guard let stmt = stmts.first else {
             let loc = SourceLocation(range: query.startIndex..<query.endIndex, line: 0, column: 0)
             diagnostics.add(.init("Query has no statements", at: loc))
-            self.diagnostics.merge(diagnostics)
-            return diagnostics
+            return (nil, diagnostics)
         }
 
         let stmtWithDef = stmt.with(
@@ -68,9 +53,7 @@ public struct Compiler {
             )
         )
 
-        self.queries.append(stmtWithDef)
-        self.diagnostics.merge(diagnostics)
-        return diagnostics
+        return (stmtWithDef, diagnostics)
     }
 
     mutating func compile<Validator>(

Sources/Compiler/Driver.swift

@@ -118,7 +118,7 @@ public actor Driver {
         var compiler = Compiler()
         compiler.schema = currentSchema
 
-        let diagnostics = switch usage {
+        let (statements, diagnostics) = switch usage {
         case .migration:
             compiler.compile(migration: fileContents)
         case .queries:
@@ -131,7 +131,7 @@ public actor Driver {
             fileName: file,
             usage: usage,
             diagnostics: diagnostics,
-            statements: compiler.queries,
+            statements: statements,
             schema: compiler.schema
         )
 

Sources/Compiler/Sema/StmtTypeChecker.swift

@@ -961,7 +961,8 @@ extension StmtTypeChecker {
             for (name, def) in columnsDefs {
                 columns[name.value] = typeFor(
                     column: def,
-                    tableColumns: columns
+                    tableColumns: columns,
+                    tableName: createTable.name.value
                 )
             }
 
@@ -1001,7 +1002,11 @@ extension StmtTypeChecker {
                 newColumns[column.key == oldName.value ? newName.value : column.key] = column.value
             }
         case let .addColumn(column):
-            table.columns[column.name.value] = typeFor(column: column, tableColumns: table.columns)
+            table.columns[column.name.value] = typeFor(
+                column: column,
+                tableColumns: table.columns,
+                tableName: table.name
+            )
         case let .dropColumn(column):
             table.columns[column.value] = nil
         }
@@ -1040,7 +1045,8 @@ extension StmtTypeChecker {
     /// Will figure out the final SQL column type from the syntax
     private mutating func typeFor(
         column: borrowing ColumnDefSyntax,
-        tableColumns: borrowing Columns
+        tableColumns: borrowing Columns,
+        tableName: Substring
     ) -> Type {
         var isNotNullable = false
         for constraint in column.constraints {
@@ -1059,7 +1065,23 @@ extension StmtTypeChecker {
                     _ = typeChecker.typeCheck(expr)
                 }
             case .foreignKey(let fk):
-                if schema[fk.foreignTable.value] == nil {
+                if fk.foreignTable.value == tableName {
+                    for foreignColumn in fk.foreignColumns {
+                        // Column constraints can reference the column they are
+                        // declared for so if its this table and this column then ignore it.
+                        guard column.name.value != foreignColumn.value else { continue }
+                        
+                        if tableColumns[foreignColumn.value] == nil {
+                            diagnostics.add(.columnDoesNotExist(foreignColumn))
+                        }
+                    }
+                } else if let table = schema[fk.foreignTable.value] {
+                    for foreignColumn in fk.foreignColumns {
+                        if table.columns[foreignColumn.value] == nil {
+                            diagnostics.add(.columnDoesNotExist(foreignColumn))
+                        }
+                    }
+                } else {
                     diagnostics.add(.tableDoesNotExist(fk.foreignTable))
                 }
             case .generated(let expr, _):
@@ -1095,6 +1117,21 @@ extension StmtTypeChecker {
         }
     }
 
+    private mutating func typeCheck(
+        fk: ForeignKeyClauseSyntax,
+        column: ColumnDefSyntax,
+        tableColumns: Columns
+    ) {
+        for foreignColumn in fk.foreignColumns {
+            // Column constraints can oddly reference themselves
+            guard column.name.value != foreignColumn.value else { continue }
+            
+            if tableColumns[foreignColumn.value] == nil {
+                diagnostics.add(.columnDoesNotExist(foreignColumn))
+            }
+        }
+    }
+    
     /// Gets the column names of the primary key and validates them
     private mutating func primaryKey(
         of stmt: CreateTableStmtSyntax,

Sources/FeatherMacros/DatabaseMacro.swift

@@ -34,28 +34,37 @@ extension DatabaseMacro: MemberMacro {
         }
 
         var compiler = Compiler()
+        var queries: [Statement] = []
 
         for (migration, expr) in migrations {
-            for diag in compiler.compile(migration: migration) {
+            let (_, diagnostics) = compiler.compile(migration: migration)
+            
+            for diag in diagnostics {
                 context.addDiagnostics(from: diag, node: expr)
             }
         }
 
         for (name, variable) in variables {
             guard let queryMacro = variable.queryMacroInputsIfIsQuery(in: context) else { continue }
 
-            for diag in compiler.compile(
+            let (statement, diagnostics) = compiler.compile(
                 query: queryMacro.source,
                 named: name.removingQuerySuffix(),
                 inputType: queryMacro.inputName,
                 outputType: queryMacro.outputName
-            ) {
+            )
+            
+            for diag in diagnostics {
                 context.addDiagnostics(from: diag, node: variable)
             }
+            
+            if let statement {
+                queries.append(statement)
+            }
         }
 
         let (generatedTables, generatedQueries) = try SwiftLanguage.assemble(
-            queries: [(nil, compiler.queries)],
+            queries: [(nil, queries)],
             schema: compiler.schema
         )
 

Tests/CompilerTests/Compiler/CompileCreateTable.sql

@@ -60,7 +60,7 @@ CREATE TABLE baz (
 -- CHECK:   KIND normal
 -- CHECK-ERROR: Table 'qux' already has a primary key
 CREATE TABLE qux (
-    foo TEXT PRIMARY KEY,
+    foo TEXT PRIMARY KEY ON CONFLICT REPLACE REFERENCES qux (foo) ON DELETE CASCADE,
     bar INTEGER,
     PRIMARY KEY (bar)
 ) STRICT;

Tests/CompilerTests/CompilerTests.swift

@@ -57,8 +57,8 @@ class CompilerTests: XCTestCase {
             sqlFile: "CompileIsSingleResult",
             parse: { contents in
                 var compiler = Compiler()
-                _ = compiler.compile(queries: contents)
-                return compiler.queries
+                let (statements, _) = compiler.compile(queries: contents)
+                return statements
                     .filter{ !($0.syntax is CreateTableStmtSyntax) }
                     .map { $0.outputCardinality.rawValue.uppercased() }
             }