126 lines
10 KiB
Swift
126 lines
10 KiB
Swift
import Foundation
|
|
import SQLite3
|
|
|
|
private let rdaiSQLiteTransient = unsafeBitCast(-1, to: sqlite3_destructor_type.self)
|
|
|
|
/// Durable local store. Each resource replacement is atomic, so interrupted
|
|
/// indexing can safely resume from the first resource without a checkpoint.
|
|
public actor RDAISQLiteIndexStore {
|
|
private var database: OpaquePointer?
|
|
private let encoder = JSONEncoder()
|
|
private let decoder = JSONDecoder()
|
|
|
|
public init(directory: URL? = nil) throws {
|
|
let base = directory ?? FileManager.default.urls(for: .applicationSupportDirectory, in: .userDomainMask).first!
|
|
try FileManager.default.createDirectory(at: base, withIntermediateDirectories: true)
|
|
let url = base.appendingPathComponent("RDAIReaderView.sqlite")
|
|
guard sqlite3_open_v2(url.path, &database, SQLITE_OPEN_CREATE | SQLITE_OPEN_READWRITE | SQLITE_OPEN_FULLMUTEX, nil) == SQLITE_OK else {
|
|
throw RDAIError.storageFailure("open")
|
|
}
|
|
try Self.configure(database: database)
|
|
}
|
|
|
|
deinit { sqlite3_close(database) }
|
|
|
|
public func replace(document: RDAIDocumentDescriptor, resource: RDAIResourceDescriptor, analysis: RDAIAnalysisResult, sourceHash: String) throws {
|
|
try transaction {
|
|
try execute("INSERT INTO documents(id,title,format,revision) VALUES(?,?,?,?) ON CONFLICT(id) DO UPDATE SET title=excluded.title,format=excluded.format,revision=excluded.revision", [document.identifier.rawValue, document.title, document.format.rawValue, document.contentRevision])
|
|
try execute("DELETE FROM resources WHERE document_id=? AND resource_id=?", [document.identifier.rawValue, resource.identifier.rawValue])
|
|
try execute("INSERT INTO resources(document_id,resource_id,source_hash,payload) VALUES(?,?,?,?)", [document.identifier.rawValue, resource.identifier.rawValue, sourceHash, try json(resource)])
|
|
try execute("DELETE FROM passages WHERE document_id=? AND resource_id=?", [document.identifier.rawValue, resource.identifier.rawValue])
|
|
try execute("DELETE FROM passage_fts WHERE document_id=? AND resource_id=?", [document.identifier.rawValue, resource.identifier.rawValue])
|
|
try execute("DELETE FROM entity_mentions WHERE document_id=? AND resource_id=?", [document.identifier.rawValue, resource.identifier.rawValue])
|
|
for passage in analysis.passages {
|
|
try execute("INSERT INTO passages(id,document_id,resource_id,content_hash,payload) VALUES(?,?,?,?,?)", [passage.id, document.identifier.rawValue, resource.identifier.rawValue, passage.contentHash, try json(passage)])
|
|
try execute("INSERT INTO passage_fts(id,document_id,resource_id,text) VALUES(?,?,?,?)", [passage.id, document.identifier.rawValue, resource.identifier.rawValue, passage.text])
|
|
}
|
|
for mention in analysis.entityMentions {
|
|
try execute("INSERT INTO entity_mentions(id,document_id,resource_id,normalized_name,payload) VALUES(?,?,?,?,?)", [mention.id, document.identifier.rawValue, resource.identifier.rawValue, mention.normalizedName, try json(mention)])
|
|
}
|
|
}
|
|
}
|
|
|
|
public func resourceHash(documentID: RDAIDocumentIdentifier, resourceID: RDAIResourceIdentifier) throws -> String? {
|
|
try query("SELECT source_hash FROM resources WHERE document_id=? AND resource_id=?", [documentID.rawValue, resourceID.rawValue]).first
|
|
}
|
|
|
|
public func passages(documentID: RDAIDocumentIdentifier) throws -> [RDAIPassage] {
|
|
try query("SELECT payload FROM passages WHERE document_id=? ORDER BY rowid", [documentID.rawValue]).compactMap { try? decode(RDAIPassage.self, $0) }
|
|
}
|
|
|
|
public func entities(documentID: RDAIDocumentIdentifier) throws -> [RDAIEntityMention] {
|
|
try query("SELECT payload FROM entity_mentions WHERE document_id=? ORDER BY rowid", [documentID.rawValue]).compactMap { try? decode(RDAIEntityMention.self, $0) }
|
|
}
|
|
|
|
public func resourceOrder(documentID: RDAIDocumentIdentifier, resourceID: RDAIResourceIdentifier) throws -> Int? {
|
|
guard let payload = try query("SELECT payload FROM resources WHERE document_id=? AND resource_id=?", [documentID.rawValue, resourceID.rawValue]).first,
|
|
let resource = try? decode(RDAIResourceDescriptor.self, payload) else { return nil }
|
|
return resource.order
|
|
}
|
|
|
|
public func remove(documentID: RDAIDocumentIdentifier) throws {
|
|
try transaction {
|
|
try execute("DELETE FROM passage_fts WHERE document_id=?", [documentID.rawValue])
|
|
try execute("DELETE FROM documents WHERE id=?", [documentID.rawValue])
|
|
}
|
|
}
|
|
public func removeAll() throws {
|
|
try transaction {
|
|
try execute("DELETE FROM documents")
|
|
try execute("DELETE FROM passage_fts")
|
|
}
|
|
}
|
|
|
|
public func saveCheckpoint(documentID: RDAIDocumentIdentifier, resourceOrder: Int) throws {
|
|
try execute("INSERT INTO index_jobs(document_id,completed_order,updated_at) VALUES(?,?,?) ON CONFLICT(document_id) DO UPDATE SET completed_order=excluded.completed_order,updated_at=excluded.updated_at", [documentID.rawValue, String(resourceOrder), String(Date().timeIntervalSince1970)])
|
|
}
|
|
|
|
public func checkpoint(documentID: RDAIDocumentIdentifier) throws -> Int? {
|
|
try query("SELECT completed_order FROM index_jobs WHERE document_id=?", [documentID.rawValue]).first.flatMap(Int.init)
|
|
}
|
|
|
|
public func saveArtifact(documentID: RDAIDocumentIdentifier, key: String, payload: Data) throws {
|
|
try execute("INSERT INTO artifacts(document_id,cache_key,payload,created_at) VALUES(?,?,?,?) ON CONFLICT(document_id,cache_key) DO UPDATE SET payload=excluded.payload,created_at=excluded.created_at", [documentID.rawValue, key, payload, String(Date().timeIntervalSince1970)])
|
|
}
|
|
|
|
public func artifact(documentID: RDAIDocumentIdentifier, key: String) throws -> Data? {
|
|
try queryData("SELECT payload FROM artifacts WHERE document_id=? AND cache_key=?", [documentID.rawValue, key]).first
|
|
}
|
|
|
|
public func storageBytes() -> Int64 {
|
|
guard let database,
|
|
let path = sqlite3_db_filename(database, "main"),
|
|
let attributes = try? FileManager.default.attributesOfItem(atPath: String(cString: path)),
|
|
let size = attributes[.size] as? NSNumber else { return 0 }
|
|
return size.int64Value
|
|
}
|
|
|
|
private static func configure(database: OpaquePointer?) throws {
|
|
try executeScript(database: database, sql: """
|
|
CREATE TABLE IF NOT EXISTS schema_metadata(version INTEGER NOT NULL);
|
|
INSERT INTO schema_metadata(version) SELECT 1 WHERE NOT EXISTS(SELECT 1 FROM schema_metadata);
|
|
CREATE TABLE IF NOT EXISTS documents(id TEXT PRIMARY KEY,title TEXT NOT NULL,format TEXT NOT NULL,revision TEXT NOT NULL);
|
|
CREATE TABLE IF NOT EXISTS resources(document_id TEXT NOT NULL REFERENCES documents(id) ON DELETE CASCADE,resource_id TEXT NOT NULL,source_hash TEXT NOT NULL,payload BLOB NOT NULL,PRIMARY KEY(document_id,resource_id));
|
|
CREATE TABLE IF NOT EXISTS passages(id TEXT PRIMARY KEY,document_id TEXT NOT NULL REFERENCES documents(id) ON DELETE CASCADE,resource_id TEXT NOT NULL,content_hash TEXT NOT NULL,payload BLOB NOT NULL);
|
|
CREATE VIRTUAL TABLE IF NOT EXISTS passage_fts USING fts5(id UNINDEXED,document_id UNINDEXED,resource_id UNINDEXED,text);
|
|
CREATE TABLE IF NOT EXISTS entity_mentions(id TEXT PRIMARY KEY,document_id TEXT NOT NULL REFERENCES documents(id) ON DELETE CASCADE,resource_id TEXT NOT NULL,normalized_name TEXT NOT NULL,payload BLOB NOT NULL);
|
|
CREATE TABLE IF NOT EXISTS index_jobs(document_id TEXT PRIMARY KEY REFERENCES documents(id) ON DELETE CASCADE,completed_order INTEGER NOT NULL,updated_at REAL NOT NULL);
|
|
CREATE TABLE IF NOT EXISTS artifacts(document_id TEXT NOT NULL REFERENCES documents(id) ON DELETE CASCADE,cache_key TEXT NOT NULL,payload BLOB NOT NULL,created_at REAL NOT NULL,PRIMARY KEY(document_id,cache_key));
|
|
""")
|
|
}
|
|
|
|
private func transaction(_ body: () throws -> Void) throws { try execute("BEGIN IMMEDIATE"); do { try body(); try execute("COMMIT") } catch { try? execute("ROLLBACK"); throw error } }
|
|
private static func executeScript(database: OpaquePointer?, sql: String) throws {
|
|
var errorMessage: UnsafeMutablePointer<Int8>?
|
|
guard sqlite3_exec(database, sql, nil, nil, &errorMessage) == SQLITE_OK else {
|
|
defer { sqlite3_free(errorMessage) }
|
|
throw RDAIError.storageFailure("migration")
|
|
}
|
|
}
|
|
private func json<T: Encodable>(_ value: T) throws -> Data { try encoder.encode(value) }
|
|
private func decode<T: Decodable>(_ type: T.Type, _ value: String) throws -> T { try decoder.decode(T.self, from: Data(value.utf8)) }
|
|
private func queryData(_ sql: String, _ bindings: [String]) throws -> [Data] { var statement: OpaquePointer?; guard sqlite3_prepare_v2(database, sql, -1, &statement, nil) == SQLITE_OK else { throw RDAIError.storageFailure("prepare") }; defer { sqlite3_finalize(statement) }; for (index, value) in bindings.enumerated() { sqlite3_bind_text(statement, Int32(index + 1), value, -1, rdaiSQLiteTransient) }; var rows: [Data] = []; while sqlite3_step(statement) == SQLITE_ROW { let count = sqlite3_column_bytes(statement, 0); if let value = sqlite3_column_blob(statement, 0), count > 0 { rows.append(Data(bytes: value, count: Int(count))) } }; return rows }
|
|
private func query(_ sql: String, _ bindings: [String]) throws -> [String] { var statement: OpaquePointer?; guard sqlite3_prepare_v2(database, sql, -1, &statement, nil) == SQLITE_OK else { throw RDAIError.storageFailure("prepare") }; defer { sqlite3_finalize(statement) }; for (index, value) in bindings.enumerated() { sqlite3_bind_text(statement, Int32(index + 1), value, -1, rdaiSQLiteTransient) }; var rows: [String] = []; while sqlite3_step(statement) == SQLITE_ROW { if let value = sqlite3_column_text(statement, 0) { rows.append(String(cString: value)) } }; return rows }
|
|
private func execute(_ sql: String, _ bindings: [Any] = []) throws { var statement: OpaquePointer?; guard sqlite3_prepare_v2(database, sql, -1, &statement, nil) == SQLITE_OK else { throw RDAIError.storageFailure("prepare") }; defer { sqlite3_finalize(statement) }; for (index, value) in bindings.enumerated() { if let text = value as? String { sqlite3_bind_text(statement, Int32(index + 1), text, -1, rdaiSQLiteTransient) } else if let data = value as? Data { _ = data.withUnsafeBytes { sqlite3_bind_blob(statement, Int32(index + 1), $0.baseAddress, Int32(data.count), rdaiSQLiteTransient) } } }; guard sqlite3_step(statement) == SQLITE_DONE else { throw RDAIError.storageFailure("execute") } }
|
|
}
|