Skip to content

Commit 8981a23

Browse files
rausnitzfabianfett
andauthored
Add support for int4range, int8range, int4range[], int8range[] (vapor#330)
Co-authored-by: Fabian Fett <fabianfett@apple.com>
1 parent a290e4e commit 8981a23

File tree

7 files changed

+684
-1
lines changed

7 files changed

+684
-1
lines changed

Sources/PostgresNIO/Data/PostgresDataType.swift

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -115,6 +115,14 @@ public struct PostgresDataType: RawRepresentable, Sendable, Hashable, CustomStri
115115
public static let jsonb = PostgresDataType(3802)
116116
/// `3807` _jsonb
117117
public static let jsonbArray = PostgresDataType(3807)
118+
/// `3904`
119+
public static let int4Range = PostgresDataType(3904)
120+
/// `3905` _int4range
121+
public static let int4RangeArray = PostgresDataType(3905)
122+
/// `3926`
123+
public static let int8Range = PostgresDataType(3926)
124+
/// `3927` _int8range
125+
public static let int8RangeArray = PostgresDataType(3927)
118126

119127
/// The raw data type code recognized by PostgreSQL.
120128
public var rawValue: UInt32
@@ -180,6 +188,10 @@ public struct PostgresDataType: RawRepresentable, Sendable, Hashable, CustomStri
180188
case .uuidArray: return "UUID[]"
181189
case .jsonb: return "JSONB"
182190
case .jsonbArray: return "JSONB[]"
191+
case .int4Range: return "INT4RANGE"
192+
case .int4RangeArray: return "INT4RANGE[]"
193+
case .int8Range: return "INT8RANGE"
194+
case .int8RangeArray: return "INT8RANGE[]"
183195
default: return nil
184196
}
185197
}
@@ -201,6 +213,8 @@ public struct PostgresDataType: RawRepresentable, Sendable, Hashable, CustomStri
201213
case .jsonb: return .jsonbArray
202214
case .text: return .textArray
203215
case .varchar: return .varcharArray
216+
case .int4Range: return .int4RangeArray
217+
case .int8Range: return .int8RangeArray
204218
default: return nil
205219
}
206220
}
@@ -223,6 +237,19 @@ public struct PostgresDataType: RawRepresentable, Sendable, Hashable, CustomStri
223237
case .jsonbArray: return .jsonb
224238
case .textArray: return .text
225239
case .varcharArray: return .varchar
240+
case .int4RangeArray: return .int4Range
241+
case .int8RangeArray: return .int8Range
242+
default: return nil
243+
}
244+
}
245+
246+
/// Returns the bound type for this type if one is known.
247+
/// Returns nil if this is not a range type.
248+
@usableFromInline
249+
internal var boundType: PostgresDataType? {
250+
switch self {
251+
case .int4Range: return .int4
252+
case .int8Range: return .int8
226253
default: return nil
227254
}
228255
}

Sources/PostgresNIO/New/Data/Array+PostgresCodable.swift

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,18 @@ extension UUID: PostgresArrayEncodable {
8585
public static var psqlArrayType: PostgresDataType { .uuidArray }
8686
}
8787

88+
extension Range: PostgresArrayDecodable where Bound: PostgresRangeArrayDecodable {}
89+
90+
extension Range: PostgresArrayEncodable where Bound: PostgresRangeArrayEncodable {
91+
public static var psqlArrayType: PostgresDataType { Bound.psqlRangeArrayType }
92+
}
93+
94+
extension ClosedRange: PostgresArrayDecodable where Bound: PostgresRangeArrayDecodable {}
95+
96+
extension ClosedRange: PostgresArrayEncodable where Bound: PostgresRangeArrayEncodable {
97+
public static var psqlArrayType: PostgresDataType { Bound.psqlRangeArrayType }
98+
}
99+
88100
// MARK: Array conformances
89101

90102
extension Array: PostgresEncodable where Element: PostgresArrayEncodable {
Lines changed: 307 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,307 @@
1+
import NIOCore
2+
3+
// MARK: Protocols
4+
5+
/// A type that can be encoded into a Postgres range type where it is the bound type
6+
public protocol PostgresRangeEncodable: PostgresNonThrowingEncodable {
7+
static var psqlRangeType: PostgresDataType { get }
8+
}
9+
10+
/// A type that can be decoded into a Swift RangeExpression type from a Postgres range where it is the bound type
11+
public protocol PostgresRangeDecodable: PostgresDecodable {
12+
/// If a Postgres range type has a well-defined step,
13+
/// Postgres automatically converts it to a canonical form.
14+
/// Types such as `int4range` get converted to upper-bound-exclusive.
15+
/// This method is needed when converting an upper bound to inclusive.
16+
/// It should throw if the type lacks a well-defined step.
17+
func upperBoundExclusiveToUpperBoundInclusive() throws -> Self
18+
19+
/// Postgres does not store any bound values for empty ranges,
20+
/// but Swift requires a value to initialize an empty Range<Bound>.
21+
static var valueForEmptyRange: Self { get }
22+
}
23+
24+
/// A type that can be encoded into a Postgres range array type where it is the bound type
25+
public protocol PostgresRangeArrayEncodable: PostgresRangeEncodable {
26+
static var psqlRangeArrayType: PostgresDataType { get }
27+
}
28+
29+
/// A type that can be decoded into a Swift RangeExpression array type from a Postgres range array where it is the bound type
30+
public protocol PostgresRangeArrayDecodable: PostgresRangeDecodable {}
31+
32+
// MARK: Bound conformances
33+
34+
extension FixedWidthInteger where Self: PostgresRangeDecodable {
35+
public func upperBoundExclusiveToUpperBoundInclusive() -> Self {
36+
return self - 1
37+
}
38+
39+
public static var valueForEmptyRange: Self {
40+
return .zero
41+
}
42+
}
43+
44+
extension Int32: PostgresRangeEncodable {
45+
public static var psqlRangeType: PostgresDataType { return .int4Range }
46+
}
47+
48+
extension Int32: PostgresRangeDecodable {}
49+
50+
extension Int32: PostgresRangeArrayEncodable {
51+
public static var psqlRangeArrayType: PostgresDataType { return .int4RangeArray }
52+
}
53+
54+
extension Int32: PostgresRangeArrayDecodable {}
55+
56+
extension Int64: PostgresRangeEncodable {
57+
public static var psqlRangeType: PostgresDataType { return .int8Range }
58+
}
59+
60+
extension Int64: PostgresRangeDecodable {}
61+
62+
extension Int64: PostgresRangeArrayEncodable {
63+
public static var psqlRangeArrayType: PostgresDataType { return .int8RangeArray }
64+
}
65+
66+
extension Int64: PostgresRangeArrayDecodable {}
67+
68+
// MARK: PostgresRange
69+
70+
@usableFromInline
71+
struct PostgresRange<B> {
72+
@usableFromInline let lowerBound: B?
73+
@usableFromInline let upperBound: B?
74+
@usableFromInline let isLowerBoundInclusive: Bool
75+
@usableFromInline let isUpperBoundInclusive: Bool
76+
77+
@inlinable
78+
init(
79+
lowerBound: B?,
80+
upperBound: B?,
81+
isLowerBoundInclusive: Bool,
82+
isUpperBoundInclusive: Bool
83+
) {
84+
self.lowerBound = lowerBound
85+
self.upperBound = upperBound
86+
self.isLowerBoundInclusive = isLowerBoundInclusive
87+
self.isUpperBoundInclusive = isUpperBoundInclusive
88+
}
89+
}
90+
91+
/// Used by Postgres to represent certain range properties
92+
@usableFromInline
93+
struct PostgresRangeFlag {
94+
@usableFromInline static let isEmpty: UInt8 = 0x01
95+
@usableFromInline static let isLowerBoundInclusive: UInt8 = 0x02
96+
@usableFromInline static let isUpperBoundInclusive: UInt8 = 0x04
97+
}
98+
99+
extension PostgresRange: PostgresDecodable where B: PostgresRangeDecodable {
100+
@inlinable
101+
init<JSONDecoder: PostgresJSONDecoder>(
102+
from byteBuffer: inout ByteBuffer,
103+
type: PostgresDataType,
104+
format: PostgresFormat,
105+
context: PostgresDecodingContext<JSONDecoder>
106+
) throws {
107+
guard case .binary = format else {
108+
throw PostgresDecodingError.Code.failure
109+
}
110+
111+
guard let boundType: PostgresDataType = type.boundType else {
112+
throw PostgresDecodingError.Code.failure
113+
}
114+
115+
// flags byte contains certain properties of the range
116+
guard let flags: UInt8 = byteBuffer.readInteger(as: UInt8.self) else {
117+
throw PostgresDecodingError.Code.failure
118+
}
119+
120+
let isEmpty: Bool = flags & PostgresRangeFlag.isEmpty != 0
121+
if isEmpty {
122+
self = PostgresRange<B>(
123+
lowerBound: B.valueForEmptyRange,
124+
upperBound: B.valueForEmptyRange,
125+
isLowerBoundInclusive: true,
126+
isUpperBoundInclusive: false
127+
)
128+
return
129+
}
130+
131+
guard let lowerBoundSize: Int32 = byteBuffer.readInteger(as: Int32.self),
132+
Int(lowerBoundSize) == MemoryLayout<B>.size,
133+
var lowerBoundBytes: ByteBuffer = byteBuffer.readSlice(length: Int(lowerBoundSize))
134+
else {
135+
throw PostgresDecodingError.Code.failure
136+
}
137+
138+
let lowerBound: B = try B(from: &lowerBoundBytes, type: boundType, format: format, context: context)
139+
140+
guard let upperBoundSize = byteBuffer.readInteger(as: Int32.self),
141+
Int(upperBoundSize) == MemoryLayout<B>.size,
142+
var upperBoundBytes: ByteBuffer = byteBuffer.readSlice(length: Int(upperBoundSize))
143+
else {
144+
throw PostgresDecodingError.Code.failure
145+
}
146+
147+
let upperBound: B = try B(from: &upperBoundBytes, type: boundType, format: format, context: context)
148+
149+
let isLowerBoundInclusive: Bool = flags & PostgresRangeFlag.isLowerBoundInclusive != 0
150+
let isUpperBoundInclusive: Bool = flags & PostgresRangeFlag.isUpperBoundInclusive != 0
151+
152+
self = PostgresRange<B>(
153+
lowerBound: lowerBound,
154+
upperBound: upperBound,
155+
isLowerBoundInclusive: isLowerBoundInclusive,
156+
isUpperBoundInclusive: isUpperBoundInclusive
157+
)
158+
159+
}
160+
}
161+
162+
extension PostgresRange: PostgresEncodable & PostgresNonThrowingEncodable where B: PostgresRangeEncodable {
163+
@usableFromInline
164+
static var psqlType: PostgresDataType { return B.psqlRangeType }
165+
166+
@usableFromInline
167+
static var psqlFormat: PostgresFormat { return .binary }
168+
169+
@inlinable
170+
func encode<JSONEncoder: PostgresJSONEncoder>(into byteBuffer: inout ByteBuffer, context: PostgresEncodingContext<JSONEncoder>) {
171+
// flags byte contains certain properties of the range
172+
var flags: UInt8 = 0
173+
if self.isLowerBoundInclusive {
174+
flags |= PostgresRangeFlag.isLowerBoundInclusive
175+
}
176+
if self.isUpperBoundInclusive {
177+
flags |= PostgresRangeFlag.isUpperBoundInclusive
178+
}
179+
180+
let boundMemorySize = Int32(MemoryLayout<B>.size)
181+
182+
byteBuffer.writeInteger(flags)
183+
if let lowerBound: B = self.lowerBound {
184+
byteBuffer.writeInteger(boundMemorySize)
185+
lowerBound.encode(into: &byteBuffer, context: context)
186+
}
187+
if let upperBound: B = self.upperBound {
188+
byteBuffer.writeInteger(boundMemorySize)
189+
upperBound.encode(into: &byteBuffer, context: context)
190+
}
191+
}
192+
}
193+
194+
extension PostgresRange where B: Comparable {
195+
@inlinable
196+
init(range: Range<B>) {
197+
self.lowerBound = range.lowerBound
198+
self.upperBound = range.upperBound
199+
self.isLowerBoundInclusive = true
200+
self.isUpperBoundInclusive = false
201+
}
202+
203+
@inlinable
204+
init(closedRange: ClosedRange<B>) {
205+
self.lowerBound = closedRange.lowerBound
206+
self.upperBound = closedRange.upperBound
207+
self.isLowerBoundInclusive = true
208+
self.isUpperBoundInclusive = true
209+
}
210+
}
211+
212+
// MARK: Range
213+
214+
extension Range: PostgresEncodable where Bound: PostgresRangeEncodable {
215+
public static var psqlType: PostgresDataType { return Bound.psqlRangeType }
216+
public static var psqlFormat: PostgresFormat { return .binary }
217+
218+
@inlinable
219+
public func encode<JSONEncoder: PostgresJSONEncoder>(
220+
into byteBuffer: inout ByteBuffer,
221+
context: PostgresEncodingContext<JSONEncoder>
222+
) {
223+
let postgresRange = PostgresRange<Bound>(range: self)
224+
postgresRange.encode(into: &byteBuffer, context: context)
225+
}
226+
}
227+
228+
extension Range: PostgresNonThrowingEncodable where Bound: PostgresRangeEncodable {}
229+
230+
extension Range: PostgresDecodable where Bound: PostgresRangeDecodable {
231+
@inlinable
232+
public init<JSONDecoder: PostgresJSONDecoder>(
233+
from buffer: inout ByteBuffer,
234+
type: PostgresDataType,
235+
format: PostgresFormat,
236+
context: PostgresDecodingContext<JSONDecoder>
237+
) throws {
238+
let postgresRange = try PostgresRange<Bound>(
239+
from: &buffer,
240+
type: type,
241+
format: format,
242+
context: context
243+
)
244+
245+
guard let lowerBound: Bound = postgresRange.lowerBound,
246+
let upperBound: Bound = postgresRange.upperBound,
247+
postgresRange.isLowerBoundInclusive,
248+
!postgresRange.isUpperBoundInclusive
249+
else {
250+
throw PostgresDecodingError.Code.failure
251+
}
252+
253+
self = lowerBound..<upperBound
254+
}
255+
}
256+
257+
// MARK: ClosedRange
258+
259+
extension ClosedRange: PostgresEncodable where Bound: PostgresRangeEncodable {
260+
public static var psqlType: PostgresDataType { return Bound.psqlRangeType }
261+
public static var psqlFormat: PostgresFormat { return .binary }
262+
263+
@inlinable
264+
public func encode<JSONEncoder: PostgresJSONEncoder>(
265+
into byteBuffer: inout ByteBuffer,
266+
context: PostgresEncodingContext<JSONEncoder>
267+
) {
268+
let postgresRange = PostgresRange<Bound>(closedRange: self)
269+
postgresRange.encode(into: &byteBuffer, context: context)
270+
}
271+
}
272+
273+
extension ClosedRange: PostgresNonThrowingEncodable where Bound: PostgresRangeEncodable {}
274+
275+
extension ClosedRange: PostgresDecodable where Bound: PostgresRangeDecodable {
276+
@inlinable
277+
public init<JSONDecoder: PostgresJSONDecoder>(
278+
from buffer: inout ByteBuffer,
279+
type: PostgresDataType,
280+
format: PostgresFormat,
281+
context: PostgresDecodingContext<JSONDecoder>
282+
) throws {
283+
let postgresRange = try PostgresRange<Bound>(
284+
from: &buffer,
285+
type: type,
286+
format: format,
287+
context: context
288+
)
289+
290+
guard let lowerBound: Bound = postgresRange.lowerBound,
291+
var upperBound: Bound = postgresRange.upperBound,
292+
postgresRange.isLowerBoundInclusive
293+
else {
294+
throw PostgresDecodingError.Code.failure
295+
}
296+
297+
if !postgresRange.isUpperBoundInclusive {
298+
upperBound = try upperBound.upperBoundExclusiveToUpperBoundInclusive()
299+
}
300+
301+
if lowerBound > upperBound {
302+
throw PostgresDecodingError.Code.failure
303+
}
304+
305+
self = lowerBound...upperBound
306+
}
307+
}

0 commit comments

Comments
 (0)