|
| 1 | +import Async |
| 2 | +import CodableKit |
| 3 | +import FluentSQL |
| 4 | +import Foundation |
| 5 | + |
| 6 | +/// Adds ability to do basic Fluent queries using a `PostgreSQLDatabase`. |
| 7 | +extension PostgreSQLDatabase: QuerySupporting { |
| 8 | + /// See `QuerySupporting.execute` |
| 9 | + public static func execute<I, D>(query: DatabaseQuery<PostgreSQLDatabase>, into stream: I, on connection: PostgreSQLConnection) |
| 10 | + where I: Async.InputStream, D: Decodable, D == I.Input |
| 11 | + { |
| 12 | + let future = Future<Void>.flatMap { |
| 13 | + // Convert Fluent `DatabaseQuery` to generic FluentSQL `DataQuery` |
| 14 | + var (sqlQuery, bindValues) = query.makeDataQuery() |
| 15 | + |
| 16 | + // If the query has an Encodable model attached serialize it. |
| 17 | + // Dictionary keys should be added to the DataQuery as columns. |
| 18 | + // Dictionary values should be added to the parameterized array. |
| 19 | + let modelData: [PostgreSQLData] |
| 20 | + if let model = query.data { |
| 21 | + let encoder = PostgreSQLRowEncoder() |
| 22 | + try model.encode(to: encoder) |
| 23 | + sqlQuery.columns += encoder.data.keys.map { key in |
| 24 | + return DataColumn(table: query.entity, name: key) |
| 25 | + } |
| 26 | + modelData = .init(encoder.data.values) |
| 27 | + } else { |
| 28 | + modelData = [] |
| 29 | + } |
| 30 | + |
| 31 | + // Create a PostgreSQL-flavored SQL serializer to create a SQL string |
| 32 | + let sqlSerializer = PostgreSQLSQLSerializer() |
| 33 | + let sqlString = sqlSerializer.serialize(data: sqlQuery) |
| 34 | + |
| 35 | + // Combine the query data with bind values from filters. |
| 36 | + // All bind values must come _after_ the columns section of the query. |
| 37 | + let parameters = try modelData + bindValues.map { bind in |
| 38 | + let encodable = bind.encodable |
| 39 | + guard let convertible = encodable as? PostgreSQLDataCustomConvertible else { |
| 40 | + let type = Swift.type(of: encodable) |
| 41 | + throw PostgreSQLError( |
| 42 | + identifier: "convertible", |
| 43 | + reason: "Unsupported encodable type: \(type)", |
| 44 | + suggestedFixes: [ |
| 45 | + "Conform \(type) to PostgreSQLDataCustomConvertible" |
| 46 | + ], |
| 47 | + source: .capture() |
| 48 | + ) |
| 49 | + } |
| 50 | + return try convertible.convertToPostgreSQLData() |
| 51 | + } |
| 52 | + |
| 53 | + // Create a push stream to accept the psql output |
| 54 | + // FIXME: connect streams directly instead? |
| 55 | + let pushStream = PushStream<D>() |
| 56 | + pushStream.output(to: stream) |
| 57 | + |
| 58 | + // Run the query |
| 59 | + return try connection.query(sqlString, parameters) { row in |
| 60 | + do { |
| 61 | + let decoded = try D.init(from: PostgreSQLRowDecoder(row: row)) |
| 62 | + pushStream.push(decoded) |
| 63 | + } catch { |
| 64 | + pushStream.error(error) |
| 65 | + } |
| 66 | + } |
| 67 | + } |
| 68 | + |
| 69 | + /// Convert Future completion / error to stream |
| 70 | + future.do { |
| 71 | + // Query is complete |
| 72 | + stream.close() |
| 73 | + }.catch { error in |
| 74 | + // Query failed |
| 75 | + stream.error(error) |
| 76 | + stream.close() |
| 77 | + } |
| 78 | + } |
| 79 | + |
| 80 | + /// See `QuerySupporting.modelEvent` |
| 81 | + public static func modelEvent<M>(event: ModelEvent, model: M, on connection: PostgreSQLConnection) -> Future<M> |
| 82 | + where PostgreSQLDatabase == M.Database, M: Model |
| 83 | + { |
| 84 | + switch event { |
| 85 | + case .willCreate: |
| 86 | + if M.ID.self == UUID.self { |
| 87 | + var model = model |
| 88 | + model.fluentID = UUID() as? M.ID |
| 89 | + return Future(model) |
| 90 | + } |
| 91 | + case .didCreate: |
| 92 | + if M.ID.self == Int.self { |
| 93 | + return connection.simpleQuery("SELECT LASTVAL();").map(to: M.self) { row in |
| 94 | + var model = model |
| 95 | + try model.fluentID = row[0]["lastval"]?.decode(Int.self) as? M.ID |
| 96 | + return model |
| 97 | + } |
| 98 | + } |
| 99 | + default: break |
| 100 | + } |
| 101 | + |
| 102 | + return Future(model) |
| 103 | + } |
| 104 | +} |
0 commit comments