Merge pull request #7 from PureSwift/feature/function

Add custom SQL function support

Author: colemancda

Package.swift

@@ -34,7 +34,7 @@ let package = Package(
     dependencies: [
         .package(
             url: "https://github.com/PureSwift/CoreModel",
-            from: "2.7.2"
+            from: "2.8.0"
         ),
         sqliteDependency
     ],

Sources/CoreModelSQLite/AttributeValue.swift

@@ -43,6 +43,26 @@ internal extension AttributeValue {
         }
     }
 
+    /// Decode from a SQLite binding value with no declared attribute type (e.g. a
+    /// raw argument passed into a custom SQL function), inferring the value's shape
+    /// from the binding's runtime type.
+    init(binding: Binding?) {
+        switch binding {
+        case .none:
+            self = .null
+        case let value as Int64:
+            self = .int64(value)
+        case let value as Double:
+            self = .double(value)
+        case let value as String:
+            self = .string(value)
+        case let value as Blob:
+            self = .data(Data(value.bytes))
+        default:
+            self = .null
+        }
+    }
+
     /// Decode from a SQLite binding value, interpreting it according to the declared attribute type.
     init(binding: Binding?, type: AttributeType) throws {
         guard let binding else {

Sources/CoreModelSQLite/Database.swift

@@ -111,6 +111,44 @@ extension SQLiteDatabase: ModelStorage {
         try connection.delete(entity, for: ids, model: model)
         invalidateCache(for: [entity])
     }
+
+    /// Registers a custom scalar function so it can be invoked from a predicate or sort
+    /// descriptor via ``FetchRequest/Predicate/Expression/function(_:)``.
+    ///
+    /// - Important: Only supported on Apple platforms. The underlying SQLite.swift
+    ///   `createFunction` registers the callback through `@convention(block)` +
+    ///   `unsafeBitCast`, which is unreliable on non-Apple platforms and can corrupt
+    ///   SQLite's heap (upstream https://github.com/stephencelis/SQLite.swift/issues/1071).
+    public func register(function: DatabaseFunction) async throws {
+        connection.register(function: function)
+    }
+}
+
+public extension SQLiteDatabase {
+
+    /// Execute raw SQL against the underlying connection — e.g. to create and
+    /// maintain an R*Tree or other virtual table. CoreModelSQLite does not create,
+    /// sync, or otherwise know about any virtual table itself; that is entirely the
+    /// caller's responsibility.
+    func execute(_ sql: String, _ bindings: [Binding?] = []) throws {
+        try connection.run(sql, bindings)
+    }
+}
+
+internal extension SQLite.Connection {
+
+    /// Registers a `DatabaseFunction` with this connection, bridging SQLite's untyped
+    /// `Binding` values to/from `AttributeValue` at the boundary.
+    func register(function: DatabaseFunction) {
+        createFunction(
+            function.name,
+            argumentCount: function.argumentCount.map { UInt($0) },
+            deterministic: function.deterministic
+        ) { arguments in
+            let values: [AttributeValue?] = arguments.map { AttributeValue(binding: $0) }
+            return function.evaluate(values)?.binding
+        }
+    }
 }
 
 internal extension SQLiteDatabase {
@@ -352,13 +390,23 @@ internal extension FetchRequest {
             bindings += fragment.bindings
         }
         if sortDescriptors.isEmpty == false {
-            let terms = try sortDescriptors.map { sort -> String in
-                guard entity.hasColumn(for: sort.property) else {
-                    throw SQLiteDatabaseError.invalidProperty(sort.property, entity.id)
+            let placeholderPredicate = FetchRequest.Predicate.value(true)
+            let fragments = try sortDescriptors.map { sort -> SQLFragment in
+                switch sort.term {
+                case let .property(property):
+                    guard entity.hasColumn(for: property) else {
+                        throw SQLiteDatabaseError.invalidProperty(property, entity.id)
+                    }
+                    let sql = property.rawValue.quotedIdentifier + (sort.ascending ? " ASC" : " DESC")
+                    return SQLFragment(sql: sql, bindings: [])
+                case let .function(function):
+                    let functionFragment = try function.sqlFragment(for: entity, predicate: placeholderPredicate)
+                    let sql = functionFragment.sql + (sort.ascending ? " ASC" : " DESC")
+                    return SQLFragment(sql: sql, bindings: functionFragment.bindings)
                 }
-                return sort.property.rawValue.quotedIdentifier + (sort.ascending ? " ASC" : " DESC")
             }
-            sql += " ORDER BY " + terms.joined(separator: ", ")
+            sql += " ORDER BY " + fragments.map(\.sql).joined(separator: ", ")
+            bindings += fragments.flatMap(\.bindings)
         } else {
             // match CoreData's default behavior of sorting by object ID when no sort descriptors are provided
             sql += " ORDER BY \(SQLiteDatabase.primaryKeyColumn.quotedIdentifier) ASC"

Sources/CoreModelSQLite/Predicate.swift

@@ -74,6 +74,35 @@ internal extension FetchRequest.Predicate.Comparison {
         predicate: FetchRequest.Predicate
     ) throws -> SQLFragment {
 
+        // `function(...) <operator> constant` comparisons compile to a SQL function call.
+        if case let .function(function) = left {
+            guard modifier == nil else {
+                throw SQLiteDatabaseError.invalidPredicate(predicate)
+            }
+            let functionFragment = try function.sqlFragment(for: entity, predicate: predicate)
+            switch type {
+            case .lessThan, .lessThanOrEqualTo, .greaterThan, .greaterThanOrEqualTo:
+                let value = try right.constantBinding(predicate: predicate)
+                return SQLFragment(
+                    sql: "\(functionFragment.sql) \(type.rawValue) ?",
+                    bindings: functionFragment.bindings + [value]
+                )
+            case .equalTo, .notEqualTo:
+                let value = try right.constantBinding(predicate: predicate)
+                let sqlOperator = (type == .equalTo) ? "=" : "<>"
+                guard let value else {
+                    let nullOperator = (type == .equalTo) ? "IS NULL" : "IS NOT NULL"
+                    return SQLFragment(sql: "\(functionFragment.sql) \(nullOperator)", bindings: functionFragment.bindings)
+                }
+                return SQLFragment(
+                    sql: "\(functionFragment.sql) \(sqlOperator) ?",
+                    bindings: functionFragment.bindings + [value]
+                )
+            default:
+                throw SQLiteDatabaseError.invalidPredicate(predicate)
+            }
+        }
+
         // Only `keyPath <operator> constant` comparisons map directly to columns.
         guard case let .keyPath(keyPath) = left else {
             throw SQLiteDatabaseError.invalidPredicate(predicate)
@@ -207,8 +236,43 @@ private extension SQLFragment {
     }
 }
 
+internal extension FetchRequest.Predicate.FunctionExpression {
+
+    /// Translate a function call expression into a SQL function-call fragment,
+    /// e.g. `myFunction("lat", "lon", ?, ?)`.
+    func sqlFragment(
+        for entity: EntityDescription,
+        predicate: FetchRequest.Predicate
+    ) throws -> SQLFragment {
+        let argumentFragments = try arguments.map {
+            try $0.argumentSQLFragment(for: entity, predicate: predicate)
+        }
+        return SQLFragment(
+            sql: "\(name)(" + argumentFragments.map(\.sql).joined(separator: ", ") + ")",
+            bindings: argumentFragments.flatMap(\.bindings)
+        )
+    }
+}
+
 private extension FetchRequest.Predicate.Expression {
 
+    /// The expression as a SQL fragment suitable for use as a function argument
+    /// (a column reference, a constant placeholder, or a nested function call).
+    func argumentSQLFragment(
+        for entity: EntityDescription,
+        predicate: FetchRequest.Predicate
+    ) throws -> SQLFragment {
+        switch self {
+        case let .keyPath(keyPath):
+            let column = try entity.validateColumn(PropertyKey(rawValue: keyPath.rawValue), predicate: predicate)
+            return SQLFragment(sql: column, bindings: [])
+        case let .function(function):
+            return try function.sqlFragment(for: entity, predicate: predicate)
+        case .attribute, .relationship:
+            return SQLFragment(sql: "?", bindings: [try constantBinding(predicate: predicate)])
+        }
+    }
+
     /// The expression as a single constant binding.
     func constantBinding(predicate: FetchRequest.Predicate) throws -> Binding? {
         switch self {
@@ -223,7 +287,7 @@ private extension FetchRequest.Predicate.Expression {
             case .toMany:
                 throw SQLiteDatabaseError.invalidPredicate(predicate)
             }
-        case .keyPath:
+        case .keyPath, .function:
             throw SQLiteDatabaseError.invalidPredicate(predicate)
         }
     }

Sources/CoreModelSQLite/ViewContext.swift

@@ -54,4 +54,13 @@ public final class SQLiteViewContext: ViewContext {
     public func count(_ fetchRequest: FetchRequest) throws -> UInt {
         try connection.count(fetchRequest, model: model)
     }
+
+    /// Registers a custom function on this context's read-only connection.
+    ///
+    /// Function registration is per-connection: a function registered with a
+    /// paired ``SQLiteDatabase`` must also be registered here to be usable from
+    /// queries run through this view context.
+    public func register(function: DatabaseFunction) throws {
+        connection.register(function: function)
+    }
 }