Add custom SQL function support

Register custom scalar functions with a SQLite store and invoke them from
predicates and sort descriptors, plus a raw-SQL escape hatch for app-managed
virtual tables (e.g. R*Tree).

- SQLiteDatabase/SQLiteViewContext.register(function:) bridges a
  DatabaseFunction to Connection.createFunction, converting between
  AttributeValue and SQLite Binding at the boundary.
- SQLiteDatabase.execute(_:_:) runs raw SQL so apps can create and maintain
  their own virtual tables; the library never creates or syncs them itself.
- Predicate.swift compiles .function expressions to SQL function calls in
  WHERE clauses; Database.swift emits them in ORDER BY.

Author: colemancda

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,37 @@ extension SQLiteDatabase: ModelStorage {
         try connection.delete(entity, for: ids, model: model)
         invalidateCache(for: [entity])
     }
+
+    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 +383,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)
+    }
 }

Tests/CoreModelSQLiteTests/CustomFunctionTests.swift

@@ -0,0 +1,173 @@
+import Foundation
+import Testing
+import CoreModel
+import SQLite
+@testable import CoreModelSQLite
+
+/// A Haversine distance function, in meters, written directly in the test — CoreModelSQLite
+/// itself has no notion of "distance" or geo data; this exercises the generic
+/// `DatabaseFunction`/`.function` expression mechanism using a realistic example.
+private func haversineDistance(_ lat1: Double, _ lon1: Double, _ lat2: Double, _ lon2: Double) -> Double {
+    let earthRadius = 6_371_000.0
+    let dLat = (lat2 - lat1) * .pi / 180
+    let dLon = (lon2 - lon1) * .pi / 180
+    let a = sin(dLat / 2) * sin(dLat / 2)
+        + cos(lat1 * .pi / 180) * cos(lat2 * .pi / 180) * sin(dLon / 2) * sin(dLon / 2)
+    let c = 2 * atan2(a.squareRoot(), (1 - a).squareRoot())
+    return earthRadius * c
+}
+
+private let distanceFunction = DatabaseFunction(name: "distance", argumentCount: 4) { arguments in
+    guard case let .double(lat1) = arguments[0],
+          case let .double(lon1) = arguments[1],
+          case let .double(lat2) = arguments[2],
+          case let .double(lon2) = arguments[3]
+    else { return nil }
+    return .double(haversineDistance(lat1, lon1, lat2, lon2))
+}
+
+private func makeGeoDatabase() async throws -> SQLiteDatabase {
+    let model = Model(entities: [
+        EntityDescription(
+            id: "Site",
+            attributes: [
+                .init(id: "name", type: .string),
+                .init(id: "latitude", type: .double),
+                .init(id: "longitude", type: .double)
+            ],
+            relationships: []
+        )
+    ])
+    let database = try SQLiteDatabase(path: temporaryDatabasePath(named: "GeoTests"), model: model)
+    try await database.register(function: distanceFunction)
+    return database
+}
+
+/// Known coordinates, distances computed against Raleigh, NC (35.7796, -78.6382), the
+/// reference point every test below filters/sorts by.
+private let referenceLatitude = 35.7796
+private let referenceLongitude = -78.6382
+
+private let sites: [(id: ObjectID, name: String, latitude: Double, longitude: Double)] = [
+    ("raleigh", "Raleigh, NC", 35.7796, -78.6382),          // ~0 m
+    ("durham", "Durham, NC", 35.9940, -78.8986),             // ~30 km
+    ("charlotte", "Charlotte, NC", 35.2271, -80.8431),       // ~210 km
+    ("atlanta", "Atlanta, GA", 33.7490, -84.3880),           // ~530 km
+    ("nyc", "New York, NY", 40.7128, -74.0060)               // ~660 km
+]
+
+private func insertSites(_ database: SQLiteDatabase) async throws {
+    for site in sites {
+        try await database.insert(ModelData(
+            entity: "Site",
+            id: site.id,
+            attributes: [
+                "name": .string(site.name),
+                "latitude": .double(site.latitude),
+                "longitude": .double(site.longitude)
+            ]
+        ))
+    }
+}
+
+private func distanceExpression() -> FetchRequest.Predicate.Expression {
+    .function(.init(name: "distance", arguments: [
+        .keyPath("latitude"),
+        .keyPath("longitude"),
+        .attribute(.double(referenceLatitude)),
+        .attribute(.double(referenceLongitude))
+    ]))
+}
+
+/// Independently computed oracle, not exercising any library code, to check results against.
+private func expectedIDs(within radiusMeters: Double) -> Set<ObjectID> {
+    Set(sites.filter { haversineDistance($0.latitude, $0.longitude, referenceLatitude, referenceLongitude) <= radiusMeters }.map(\.id))
+}
+
+@Test func functionPredicateFiltersByRadius() async throws {
+    let database = try await makeGeoDatabase()
+    try await insertSites(database)
+
+    let radius = 250_000.0
+    let request = FetchRequest(
+        entity: "Site",
+        predicate: .comparison(.init(left: distanceExpression(), right: .attribute(.double(radius)), type: .lessThanOrEqualTo))
+    )
+    let ids = Set(try await database.fetchID(request))
+    #expect(ids == expectedIDs(within: radius))
+    #expect(ids == ["raleigh", "durham", "charlotte"])
+}
+
+@Test func functionSortOrdersByDistance() async throws {
+    let database = try await makeGeoDatabase()
+    try await insertSites(database)
+
+    let request = FetchRequest(
+        entity: "Site",
+        sortDescriptors: [.init(term: .function(.init(name: "distance", arguments: [
+            .keyPath("latitude"), .keyPath("longitude"),
+            .attribute(.double(referenceLatitude)), .attribute(.double(referenceLongitude))
+        ])), ascending: true)]
+    )
+    let ids = try await database.fetchID(request)
+    #expect(ids == ["raleigh", "durham", "charlotte", "atlanta", "nyc"])
+}
+
+@Test func functionFilterAndSortWithLimit() async throws {
+    let database = try await makeGeoDatabase()
+    try await insertSites(database)
+
+    let radius = 700_000.0
+    let request = FetchRequest(
+        entity: "Site",
+        sortDescriptors: [.init(term: .function(.init(name: "distance", arguments: [
+            .keyPath("latitude"), .keyPath("longitude"),
+            .attribute(.double(referenceLatitude)), .attribute(.double(referenceLongitude))
+        ])), ascending: true)],
+        predicate: .comparison(.init(left: distanceExpression(), right: .attribute(.double(radius)), type: .lessThanOrEqualTo)),
+        fetchLimit: 2
+    )
+    let ids = try await database.fetchID(request)
+    #expect(ids == ["raleigh", "durham"])
+}
+
+@MainActor
+@Test func functionRegisteredOnViewContext() async throws {
+    let path = temporaryDatabasePath(named: "GeoViewContextTests")
+    let model = Model(entities: [
+        EntityDescription(
+            id: "Site",
+            attributes: [
+                .init(id: "name", type: .string),
+                .init(id: "latitude", type: .double),
+                .init(id: "longitude", type: .double)
+            ],
+            relationships: []
+        )
+    ])
+    let database = try SQLiteDatabase(path: path, model: model)
+    try await database.register(function: distanceFunction)
+    try await insertSites(database)
+
+    let viewContext = try SQLiteViewContext(.uri(path), model: model)
+    try viewContext.register(function: distanceFunction)
+
+    let radius = 250_000.0
+    let request = FetchRequest(
+        entity: "Site",
+        predicate: .comparison(.init(left: distanceExpression(), right: .attribute(.double(radius)), type: .lessThanOrEqualTo))
+    )
+    let ids = Set(try viewContext.fetchID(request))
+    #expect(ids == expectedIDs(within: radius))
+}
+
+@Test func unrelatedEntitiesUnaffected() async throws {
+    // regression: entities/queries with no `.function` usage behave exactly as before
+    let database = try makeDatabase()
+    let people = (0..<3).map { index in
+        ModelData(entity: "Person", id: ObjectID(rawValue: "person\(index)"), attributes: ["name": .string("Person \(index)")])
+    }
+    try await database.insert(people)
+    let count = try await database.count(FetchRequest(entity: "Person"))
+    #expect(count == 3)
+}