327 lines
12 KiB
Swift
327 lines
12 KiB
Swift
//
|
|
// Copyright 2025 Signal Messenger, LLC
|
|
// SPDX-License-Identifier: AGPL-3.0-only
|
|
//
|
|
|
|
import Foundation
|
|
public import GRDB
|
|
|
|
/// An efficient-on-disk Key-Value Store.
|
|
///
|
|
/// This type is capable of storing boolean, integer, double, string, and
|
|
/// data values. They are stored in SQLite-native representations. It is not
|
|
/// natively capable of storing "complex" objects encoded via JSONEncoder or
|
|
/// NSKeyedArchiver; callers should convert these to Data themselves.
|
|
///
|
|
/// This type is not backwards compatible with `KeyValueStore` (except for
|
|
/// `Data` values), though `KeyValueStoreMigrator` can migrate values to the
|
|
/// new representation.
|
|
public struct NewKeyValueStore {
|
|
enum TableMetadata {
|
|
enum Columns {
|
|
static let collection = "collection"
|
|
static let key = "key"
|
|
static let value = "value"
|
|
}
|
|
static let tableName = "keyvalue"
|
|
}
|
|
|
|
static var tableName: String { TableMetadata.tableName }
|
|
static var collectionColumnName: String { TableMetadata.Columns.collection }
|
|
|
|
private let collection: String
|
|
|
|
public init(collection: String) {
|
|
self.collection = collection
|
|
}
|
|
|
|
/// Remove all values or crash if an error occurs.
|
|
public func removeAll(tx: DBWriteTransaction) {
|
|
failIfThrowsDatabaseError { () throws(GRDB.DatabaseError) in
|
|
try self.removeAllOrThrow(tx: tx)
|
|
}
|
|
}
|
|
|
|
/// Remove all values or throw the error that occurs.
|
|
public func removeAllOrThrow(tx: DBWriteTransaction) throws(GRDB.DatabaseError) {
|
|
return try withDatabaseError { try self._removeAllOrThrow(tx: tx) }
|
|
}
|
|
|
|
/// Fetch all the keys or crash if an error occurs.
|
|
public func fetchKeys(tx: DBReadTransaction) -> [String] {
|
|
return failIfThrowsDatabaseError { () throws(GRDB.DatabaseError) in
|
|
return try self.fetchKeysOrThrow(tx: tx)
|
|
}
|
|
}
|
|
|
|
/// Fetch all the keys or throw the error that occurs.
|
|
public func fetchKeysOrThrow(tx: DBReadTransaction) throws(GRDB.DatabaseError) -> [String] {
|
|
return try withDatabaseError { try self._getKeysOrThrow(tx: tx) }
|
|
}
|
|
|
|
/// Fetch a value (or nil if it doesn't exist) or crash if an error occurs.
|
|
public func fetchValue<T: KeyValueStoreValue>(_ type: T.Type, forKey key: String, tx: DBReadTransaction) -> T? {
|
|
return failIfThrowsDatabaseError { () throws(GRDB.DatabaseError) in
|
|
return try fetchValueOrThrow(T.self, forKey: key, tx: tx)
|
|
}
|
|
}
|
|
|
|
/// Fetch a value (or nil if it doesn't exist) or throw the error that occurs.
|
|
public func fetchValueOrThrow<T: KeyValueStoreValue>(_ type: T.Type, forKey key: String, tx: DBReadTransaction) throws(GRDB.DatabaseError) -> T? {
|
|
return try withDatabaseError {
|
|
return try self._getValueOrThrow(T.DatabaseType.self, key: key, tx: tx).map(T.init(keyValueStoreValue:))
|
|
}
|
|
}
|
|
|
|
/// Write/clear a value or crash if an error occurs.
|
|
public func writeValue<T: KeyValueStoreValue>(_ value: T?, forKey key: String, tx: DBWriteTransaction) {
|
|
failIfThrowsDatabaseError { () throws(GRDB.DatabaseError) in
|
|
try writeValueOrThrow(value, forKey: key, tx: tx)
|
|
}
|
|
}
|
|
|
|
/// Write/clear a value or throw the error that occurs.
|
|
public func writeValueOrThrow<T: KeyValueStoreValue>(_ value: T?, forKey key: String, tx: DBWriteTransaction) throws(GRDB.DatabaseError) {
|
|
try withDatabaseError { try self._setValueOrThrow(value?.keyValueStoreValue, key: key, tx: tx) }
|
|
}
|
|
|
|
/// Clear a value or crash if an error occurs.
|
|
public func removeValue(forKey key: String, tx: DBWriteTransaction) {
|
|
failIfThrowsDatabaseError { () throws(GRDB.DatabaseError) in
|
|
try self.removeValueOrThrow(forKey: key, tx: tx)
|
|
}
|
|
}
|
|
|
|
/// Clear a value or throw the error that occurs.
|
|
public func removeValueOrThrow(forKey key: String, tx: DBWriteTransaction) throws(GRDB.DatabaseError) {
|
|
try self.writeValueOrThrow(nil as Data?, forKey: key, tx: tx)
|
|
}
|
|
|
|
private func withDatabaseError<T>(_ block: () throws -> T) throws(GRDB.DatabaseError) -> T {
|
|
do {
|
|
return try block()
|
|
} catch {
|
|
throw error.forceCastToDatabaseError()
|
|
}
|
|
}
|
|
|
|
public static func logCollectionStatistics(tx: DBReadTransaction) {
|
|
Logger.info("KeyValueStore statistics:")
|
|
do {
|
|
let sql = """
|
|
SELECT \(TableMetadata.Columns.collection), COUNT(*)
|
|
FROM \(TableMetadata.tableName)
|
|
GROUP BY \(TableMetadata.Columns.collection)
|
|
ORDER BY COUNT(*) DESC
|
|
LIMIT 10
|
|
"""
|
|
let cursor = try Row.fetchCursor(tx.database, sql: sql)
|
|
while let row = try cursor.next() {
|
|
let collection: String = row[0]
|
|
let count: Int64 = row[1]
|
|
Logger.info("- \(collection): \(count) items")
|
|
}
|
|
} catch {
|
|
Logger.warn("\(error.grdbErrorForLogging)")
|
|
}
|
|
}
|
|
|
|
// MARK: - CRUD Methods
|
|
|
|
private func _getValueOrThrow<T: DatabaseValueConvertible>(_ type: T.Type, key: String, tx: DBReadTransaction) throws -> T? {
|
|
return try T.fetchOne(
|
|
tx.database,
|
|
sql: """
|
|
SELECT \(TableMetadata.Columns.value)
|
|
FROM \(TableMetadata.tableName)
|
|
WHERE
|
|
\(TableMetadata.Columns.key) = ?
|
|
AND \(TableMetadata.Columns.collection) == ?
|
|
""",
|
|
arguments: [key, collection],
|
|
)
|
|
}
|
|
|
|
private func _getKeysOrThrow(tx: DBReadTransaction) throws -> [String] {
|
|
return try String.fetchAll(
|
|
tx.database,
|
|
sql: """
|
|
SELECT \(TableMetadata.Columns.key)
|
|
FROM \(TableMetadata.tableName)
|
|
WHERE \(TableMetadata.Columns.collection) == ?
|
|
""",
|
|
arguments: [collection],
|
|
)
|
|
}
|
|
|
|
private func _setValueOrThrow<T: DatabaseValueConvertible>(_ value: T?, key: String, tx: DBWriteTransaction) throws {
|
|
let sql: String
|
|
let arguments: StatementArguments
|
|
if let value {
|
|
// See: https://www.sqlite.org/lang_UPSERT.html
|
|
sql = """
|
|
INSERT INTO \(TableMetadata.tableName) (
|
|
\(TableMetadata.Columns.key),
|
|
\(TableMetadata.Columns.collection),
|
|
\(TableMetadata.Columns.value)
|
|
) VALUES (?, ?, ?)
|
|
ON CONFLICT (
|
|
\(TableMetadata.Columns.key),
|
|
\(TableMetadata.Columns.collection)
|
|
) DO UPDATE
|
|
SET \(TableMetadata.Columns.value) = ?
|
|
"""
|
|
arguments = [key, collection, value, value]
|
|
} else {
|
|
// Setting to nil is a delete.
|
|
sql = """
|
|
DELETE FROM \(TableMetadata.tableName)
|
|
WHERE
|
|
\(TableMetadata.Columns.key) == ?
|
|
AND \(TableMetadata.Columns.collection) == ?
|
|
"""
|
|
arguments = [key, collection]
|
|
}
|
|
|
|
let statement = try tx.database.cachedStatement(sql: sql)
|
|
try statement.setArguments(arguments)
|
|
try statement.execute()
|
|
}
|
|
|
|
private func _removeAllOrThrow(tx: DBWriteTransaction) throws {
|
|
try tx.database.execute(
|
|
sql: """
|
|
DELETE
|
|
FROM \(TableMetadata.tableName)
|
|
WHERE \(TableMetadata.Columns.collection) == ?
|
|
""",
|
|
arguments: [collection]
|
|
)
|
|
}
|
|
}
|
|
|
|
// MARK: -
|
|
|
|
// These "implement" the on-disk representation for primitives supported by
|
|
// NewKeyValueStore. They are typically "put the value directly in SQLite",
|
|
// though there are two exceptions for things that (1) can be represented
|
|
// in both SQLite & Swift, (2) have a non-failable path between those
|
|
// representations, and (3) aren't bridged automatically. (A UInt64 can be
|
|
// stored in the database, but we need to pass it to SQLite as an Int64,
|
|
// and we need to read it back as an Int64.)
|
|
|
|
public protocol KeyValueStoreValue<DatabaseType> {
|
|
associatedtype DatabaseType where DatabaseType: DatabaseValueConvertible
|
|
var keyValueStoreValue: DatabaseType { get }
|
|
init(keyValueStoreValue: DatabaseType)
|
|
}
|
|
|
|
extension Data: KeyValueStoreValue {
|
|
public var keyValueStoreValue: Data { self }
|
|
public init(keyValueStoreValue: Data) { self = keyValueStoreValue }
|
|
}
|
|
|
|
extension String: KeyValueStoreValue {
|
|
public var keyValueStoreValue: String { self }
|
|
public init(keyValueStoreValue: String) { self = keyValueStoreValue }
|
|
}
|
|
|
|
extension Bool: KeyValueStoreValue {
|
|
public var keyValueStoreValue: Bool { self }
|
|
public init(keyValueStoreValue: Bool) { self = keyValueStoreValue }
|
|
}
|
|
|
|
extension Double: KeyValueStoreValue {
|
|
public var keyValueStoreValue: Double { self }
|
|
public init(keyValueStoreValue: Double) { self = keyValueStoreValue }
|
|
}
|
|
|
|
extension Int64: KeyValueStoreValue {
|
|
public var keyValueStoreValue: Int64 { self }
|
|
public init(keyValueStoreValue: Int64) { self = keyValueStoreValue }
|
|
}
|
|
|
|
extension UInt64: KeyValueStoreValue {
|
|
public var keyValueStoreValue: Int64 { Int64(bitPattern: self) }
|
|
public init(keyValueStoreValue: Int64) { self.init(bitPattern: keyValueStoreValue) }
|
|
}
|
|
|
|
extension Date: KeyValueStoreValue {
|
|
public var keyValueStoreValue: Double { self.timeIntervalSince1970 }
|
|
public init(keyValueStoreValue: Double) { self.init(timeIntervalSince1970: keyValueStoreValue) }
|
|
}
|
|
|
|
// MARK: -
|
|
|
|
struct KeyValueStoreMigrator {
|
|
private let collection: String
|
|
|
|
init(collection: String) {
|
|
self.collection = collection
|
|
}
|
|
|
|
/// Migrate a single NSKeyedArchiver-encoded value.
|
|
///
|
|
/// If an error occurs fetching from/writing to the database, that error is rethrown.
|
|
///
|
|
/// If an error occurs when parsing an NSKeyedArchiver-encoded value, that
|
|
/// error is not thrown, and the value is deleted. This maintains the old
|
|
/// behavior where decoding errors behave as if the value isn't set.
|
|
func migrateKey<V1Type: NSObject & NSSecureCoding, V2Type: DatabaseValueConvertible>(
|
|
_ key: String,
|
|
withValueOfType oldType: V1Type.Type,
|
|
toNewValue migrateValue: (V1Type) -> V2Type,
|
|
tx: DBWriteTransaction,
|
|
) throws {
|
|
let dataValue = try Data.fetchOne(
|
|
tx.database,
|
|
sql: "SELECT value FROM keyvalue WHERE collection = ? AND key = ?",
|
|
arguments: [collection, key],
|
|
)
|
|
guard let dataValue else {
|
|
// The representation of `nil` doesn't change.
|
|
return
|
|
}
|
|
// Dates may have been stored as NSDates or NSNumbers; this method handles
|
|
// both representations.
|
|
let isDateValue = V1Type.self == NSDate.self
|
|
let oldTypes = isDateValue ? [V1Type.self, NSNumber.self] : [oldType]
|
|
let oldValue: V1Type
|
|
switch try? NSKeyedUnarchiver.unarchivedObject(ofClasses: oldTypes, from: dataValue) {
|
|
case let _oldValue as V1Type:
|
|
oldValue = _oldValue
|
|
case let numberValue as NSNumber where isDateValue:
|
|
oldValue = NSDate(timeIntervalSince1970: numberValue.doubleValue) as! V1Type
|
|
default:
|
|
// The old KeyValueStore silently ignores malformed values. Do the same here.
|
|
Logger.error("Couldn't migrate '\(key)' in '\(collection)' because it was malformed.")
|
|
try tx.database.execute(
|
|
sql: "DELETE FROM keyvalue WHERE collection = ? AND key = ?",
|
|
arguments: [collection, key],
|
|
)
|
|
return
|
|
}
|
|
let newValue = migrateValue(oldValue)
|
|
try tx.database.execute(
|
|
sql: "UPDATE keyvalue SET value = ? WHERE collection = ? AND key = ?",
|
|
arguments: [newValue, collection, key],
|
|
)
|
|
}
|
|
|
|
func migrateString(_ key: String, tx: DBWriteTransaction) throws {
|
|
return try migrateKey(key, withValueOfType: NSString.self, toNewValue: { $0 as String }, tx: tx)
|
|
}
|
|
|
|
func migrateDate(_ key: String, tx: DBWriteTransaction) throws {
|
|
return try migrateKey(key, withValueOfType: NSDate.self, toNewValue: \.timeIntervalSince1970, tx: tx)
|
|
}
|
|
|
|
func migrateUInt32(_ key: String, tx: DBWriteTransaction) throws {
|
|
return try migrateKey(key, withValueOfType: NSNumber.self, toNewValue: { Int64($0.uint32Value) }, tx: tx)
|
|
}
|
|
|
|
func migrateBool(_ key: String, tx: DBWriteTransaction) throws {
|
|
return try migrateKey(key, withValueOfType: NSNumber.self, toNewValue: \.boolValue, tx: tx)
|
|
}
|
|
}
|