Skip to content

Commit 65b21d6

Browse files
MahdiBMMahdiBM
committed
dedicated tests for SimpleQueryStateMachine
1 parent d9fb651 commit 65b21d6

File tree

4 files changed

+293
-105
lines changed

4 files changed

+293
-105
lines changed

‎Sources/PostgresNIO/New/Connection State Machine/ConnectionStateMachine.swift

-3
Original file line numberDiff line numberDiff line change
@@ -961,7 +961,6 @@ struct ConnectionStateMachine {
961961
case .sendParseDescribeBindExecuteSync,
962962
.sendParseDescribeSync,
963963
.sendBindExecuteSync,
964-
.sendQuery,
965964
.succeedQuery,
966965
.succeedPreparedStatementCreation,
967966
.forwardRows,
@@ -1179,8 +1178,6 @@ extension ConnectionStateMachine {
11791178
return .sendParseDescribeBindExecuteSync(query)
11801179
case .sendBindExecuteSync(let executeStatement):
11811180
return .sendBindExecuteSync(executeStatement)
1182-
case .sendQuery(let query):
1183-
return .sendQuery(query)
11841181
case .failQuery(let requestContext, with: let error):
11851182
let cleanupContext = self.setErrorAndCreateCleanupContextIfNeeded(error)
11861183
return .failQuery(requestContext, with: error, cleanupContext: cleanupContext)

‎Sources/PostgresNIO/New/Connection State Machine/SimpleQueryStateMachine.swift

+15-2
Original file line numberDiff line numberDiff line change
@@ -62,15 +62,23 @@ struct SimpleQueryStateMachine {
6262
case .initialized:
6363
preconditionFailure("Start must be called immediatly after the query was created")
6464

65-
case .messagesSent(let queryContext),
66-
.rowDescriptionReceived(let queryContext, _):
65+
case .messagesSent(let queryContext):
6766
guard !self.isCancelled else {
6867
return .wait
6968
}
7069

7170
self.isCancelled = true
7271
return .failQuery(queryContext.promise, with: .queryCancelled)
7372

73+
case .rowDescriptionReceived(let queryContext, let columns):
74+
guard !self.isCancelled else {
75+
return .wait
76+
}
77+
78+
self.isCancelled = true
79+
self.state = .drain(columns)
80+
return .failQuery(queryContext.promise, with: .queryCancelled)
81+
7482
case .streaming(let columns, var streamStateMachine):
7583
precondition(!self.isCancelled)
7684
self.isCancelled = true
@@ -114,6 +122,11 @@ struct SimpleQueryStateMachine {
114122
return column
115123
}
116124

125+
guard !self.isCancelled else {
126+
self.state = .drain(rowDescription.columns)
127+
return .failQuery(queryContext.promise, with: .queryCancelled)
128+
}
129+
117130
self.avoidingStateMachineCoW { state in
118131
state = .rowDescriptionReceived(queryContext, columns)
119132
}

‎Tests/PostgresNIOTests/New/Connection State Machine/ExtendedQueryStateMachineTests.swift

-100
Original file line numberDiff line numberDiff line change
@@ -78,106 +78,6 @@ class ExtendedQueryStateMachineTests: XCTestCase {
7878
XCTAssertEqual(state.readyForQueryReceived(.idle), .fireEventReadyForQuery)
7979
}
8080

81-
func testExtendedQueryWithSimpleQueryWithoutDataRowsHappyPath() {
82-
var state = ConnectionStateMachine.readyForQuery()
83-
84-
let logger = Logger.psqlTest
85-
let promise = EmbeddedEventLoop().makePromise(of: PSQLRowStream.self)
86-
promise.fail(PSQLError.uncleanShutdown) // we don't care about the error at all.
87-
let query = "DELETE FROM table WHERE id=1"
88-
let queryContext = SimpleQueryContext(query: query, logger: logger, promise: promise)
89-
90-
XCTAssertEqual(state.enqueue(task: .simpleQuery(queryContext)), .sendQuery(query))
91-
XCTAssertEqual(state.commandCompletedReceived("DELETE 1"), .succeedQuery(promise, with: .init(value: .noRows(.tag("DELETE 1")), logger: logger)))
92-
XCTAssertEqual(state.readyForQueryReceived(.idle), .fireEventReadyForQuery)
93-
}
94-
95-
func testExtendedQueryWithSimpleQueryWithRowDescriptionWithoutDataRowsHappyPath() {
96-
var state = ConnectionStateMachine.readyForQuery()
97-
98-
let logger = Logger.psqlTest
99-
let promise = EmbeddedEventLoop().makePromise(of: PSQLRowStream.self)
100-
promise.fail(PSQLError.uncleanShutdown) // we don't care about the error at all.
101-
let nonExistentOID = 371280378
102-
let query = "SELECT * FROM pg_class WHERE oid = \(nonExistentOID)"
103-
let queryContext = SimpleQueryContext(query: query, logger: logger, promise: promise)
104-
105-
XCTAssertEqual(state.enqueue(task: .simpleQuery(queryContext)), .sendQuery(query))
106-
107-
let input: [RowDescription.Column] = [
108-
.init(name: "version", tableOID: 0, columnAttributeNumber: 0, dataType: .text, dataTypeSize: -1, dataTypeModifier: -1, format: .text)
109-
]
110-
XCTAssertEqual(state.rowDescriptionReceived(.init(columns: input)), .wait)
111-
XCTAssertEqual(state.commandCompletedReceived("DELETE 1"), .succeedQuery(promise, with: .init(value: .noRows(.tag("DELETE 1")), logger: logger)))
112-
XCTAssertEqual(state.readyForQueryReceived(.idle), .fireEventReadyForQuery)
113-
}
114-
115-
func testExtendedQueryWithSimpleQueryWithDataRowsHappyPath() {
116-
var state = ConnectionStateMachine.readyForQuery()
117-
118-
let logger = Logger.psqlTest
119-
let promise = EmbeddedEventLoop().makePromise(of: PSQLRowStream.self)
120-
promise.fail(PSQLError.uncleanShutdown) // we don't care about the error at all.
121-
let query = "SELECT version()"
122-
let queryContext = SimpleQueryContext(query: query, logger: logger, promise: promise)
123-
124-
XCTAssertEqual(state.enqueue(task: .simpleQuery(queryContext)), .sendQuery(query))
125-
126-
// We need to ensure that even though the row description from the wire says that we
127-
// will receive data in `.text` format, we will actually receive it in binary format,
128-
// since we requested it in binary with our bind message.
129-
let input: [RowDescription.Column] = [
130-
.init(name: "version", tableOID: 0, columnAttributeNumber: 0, dataType: .text, dataTypeSize: -1, dataTypeModifier: -1, format: .text)
131-
]
132-
let expected: [RowDescription.Column] = input.map {
133-
.init(name: $0.name, tableOID: $0.tableOID, columnAttributeNumber: $0.columnAttributeNumber, dataType: $0.dataType,
134-
dataTypeSize: $0.dataTypeSize, dataTypeModifier: $0.dataTypeModifier, format: .binary)
135-
}
136-
137-
XCTAssertEqual(state.rowDescriptionReceived(.init(columns: input)), .wait)
138-
let row1: DataRow = [ByteBuffer(string: "test1")]
139-
let result = QueryResult(value: .rowDescription(expected), logger: queryContext.logger)
140-
XCTAssertEqual(state.dataRowReceived(row1), .succeedQuery(promise, with: result))
141-
XCTAssertEqual(state.channelReadComplete(), .forwardRows([row1]))
142-
XCTAssertEqual(state.readEventCaught(), .wait)
143-
XCTAssertEqual(state.requestQueryRows(), .read)
144-
145-
let row2: DataRow = [ByteBuffer(string: "test2")]
146-
let row3: DataRow = [ByteBuffer(string: "test3")]
147-
let row4: DataRow = [ByteBuffer(string: "test4")]
148-
XCTAssertEqual(state.dataRowReceived(row2), .wait)
149-
XCTAssertEqual(state.dataRowReceived(row3), .wait)
150-
XCTAssertEqual(state.dataRowReceived(row4), .wait)
151-
XCTAssertEqual(state.channelReadComplete(), .forwardRows([row2, row3, row4]))
152-
XCTAssertEqual(state.requestQueryRows(), .wait)
153-
XCTAssertEqual(state.readEventCaught(), .read)
154-
155-
XCTAssertEqual(state.channelReadComplete(), .wait)
156-
XCTAssertEqual(state.readEventCaught(), .read)
157-
158-
let row5: DataRow = [ByteBuffer(string: "test5")]
159-
let row6: DataRow = [ByteBuffer(string: "test6")]
160-
XCTAssertEqual(state.dataRowReceived(row5), .wait)
161-
XCTAssertEqual(state.dataRowReceived(row6), .wait)
162-
163-
XCTAssertEqual(state.commandCompletedReceived("SELECT 2"), .forwardStreamComplete([row5, row6], commandTag: "SELECT 2"))
164-
XCTAssertEqual(state.readyForQueryReceived(.idle), .fireEventReadyForQuery)
165-
}
166-
167-
func testExtendedQueryWithSimpleQueryWithNoQuery() {
168-
var state = ConnectionStateMachine.readyForQuery()
169-
170-
let logger = Logger.psqlTest
171-
let promise = EmbeddedEventLoop().makePromise(of: PSQLRowStream.self)
172-
promise.fail(PSQLError.uncleanShutdown) // we don't care about the error at all.
173-
let query = "-- some comments"
174-
let queryContext = SimpleQueryContext(query: query, logger: logger, promise: promise)
175-
176-
XCTAssertEqual(state.enqueue(task: .simpleQuery(queryContext)), .sendQuery(query))
177-
XCTAssertEqual(state.emptyQueryResponseReceived(), .succeedQuery(promise, with: .init(value: .noRows(.emptyResponse), logger: logger)))
178-
XCTAssertEqual(state.readyForQueryReceived(.idle), .fireEventReadyForQuery)
179-
}
180-
18181
func testExtendedQueryWithNoQuery() {
18282
var state = ConnectionStateMachine.readyForQuery()
18383

0 commit comments

Comments
 (0)