use of com.google.cloud.spanner.pgadapter.wireprotocol.ExecuteMessage in project pgadapter by GoogleCloudPlatform.
the class ProtocolTest method testExecuteMessage.
@Test
public void testExecuteMessage() throws Exception {
byte[] messageMetadata = { 'E' };
String statementName = "some portal\0";
int totalRows = 99999;
byte[] length = intToBytes(4 + statementName.length() + 4);
byte[] value = Bytes.concat(messageMetadata, length, statementName.getBytes(), intToBytes(totalRows));
String expectedStatementName = "some portal";
DataInputStream inputStream = new DataInputStream(new ByteArrayInputStream(value));
ByteArrayOutputStream result = new ByteArrayOutputStream();
DataOutputStream outputStream = new DataOutputStream(result);
when(connectionHandler.getPortal(anyString())).thenReturn(intermediatePortalStatement);
when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
when(connectionMetadata.getInputStream()).thenReturn(inputStream);
when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
WireMessage message = ControlMessage.create(connectionHandler);
assertEquals(ExecuteMessage.class, message.getClass());
assertEquals(expectedStatementName, ((ExecuteMessage) message).getName());
assertEquals(totalRows, ((ExecuteMessage) message).getMaxRows());
verify(connectionHandler).getPortal("some portal");
ExecuteMessage messageSpy = (ExecuteMessage) spy(message);
doReturn(false).when(messageSpy).sendSpannerResult(anyInt(), any(IntermediatePortalStatement.class), any(QueryMode.class), anyLong());
messageSpy.send();
verify(intermediatePortalStatement).execute();
verify(messageSpy).sendSpannerResult(0, intermediatePortalStatement, QueryMode.EXTENDED, totalRows);
verify(connectionHandler).cleanUp(intermediatePortalStatement);
}
use of com.google.cloud.spanner.pgadapter.wireprotocol.ExecuteMessage in project pgadapter by GoogleCloudPlatform.
the class JdbcMockServerTest method testCursorSuccess.
@Test
public void testCursorSuccess() throws SQLException {
try (Connection connection = DriverManager.getConnection(createUrl())) {
connection.setAutoCommit(false);
try (PreparedStatement statement = connection.prepareStatement(SELECT_FIVE_ROWS.getSql())) {
// Fetch two rows at a time from the PG server.
statement.setFetchSize(2);
try (ResultSet resultSet = statement.executeQuery()) {
int index = 0;
while (resultSet.next()) {
assertEquals(++index, resultSet.getInt(1));
}
assertEquals(5, index);
}
}
connection.commit();
}
// The ExecuteSql request should only be sent once to Cloud Spanner.
assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
ExecuteSqlRequest executeRequest = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0);
assertEquals(QueryMode.NORMAL, executeRequest.getQueryMode());
assertEquals(SELECT_FIVE_ROWS.getSql(), executeRequest.getSql());
// 5. COMMIT
if (pgServer != null) {
List<DescribeMessage> describeMessages = getWireMessagesOfType(DescribeMessage.class);
assertEquals(1, describeMessages.size());
DescribeMessage describeMessage = describeMessages.get(0);
assertEquals(PreparedType.Portal, describeMessage.getType());
List<ExecuteMessage> executeMessages = getWireMessagesOfType(ExecuteMessage.class);
assertEquals(5, executeMessages.size());
assertEquals("", executeMessages.get(0).getName());
for (ExecuteMessage executeMessage : executeMessages.subList(1, executeMessages.size() - 1)) {
assertEquals(describeMessage.getName(), executeMessage.getName());
assertEquals(2, executeMessage.getMaxRows());
}
assertEquals("", executeMessages.get(executeMessages.size() - 1).getName());
List<ParseMessage> parseMessages = getWireMessagesOfType(ParseMessage.class);
assertEquals(3, parseMessages.size());
assertEquals("BEGIN", parseMessages.get(0).getStatement().getSql());
assertEquals(SELECT_FIVE_ROWS.getSql(), parseMessages.get(1).getStatement().getSql());
assertEquals("COMMIT", parseMessages.get(2).getStatement().getSql());
}
}
use of com.google.cloud.spanner.pgadapter.wireprotocol.ExecuteMessage in project pgadapter by GoogleCloudPlatform.
the class JdbcMockServerTest method testRandomResults.
@Test
public void testRandomResults() throws SQLException {
final int fetchSize = 3;
try (Connection connection = DriverManager.getConnection(createUrl())) {
connection.setAutoCommit(false);
try (PreparedStatement statement = connection.prepareStatement(SELECT_RANDOM.getSql())) {
statement.setFetchSize(fetchSize);
try (ResultSet resultSet = statement.executeQuery()) {
int rowCount = 0;
while (resultSet.next()) {
for (int col = 0; col < resultSet.getMetaData().getColumnCount(); col++) {
resultSet.getObject(col + 1);
}
rowCount++;
}
assertEquals(RANDOM_RESULTS_ROW_COUNT, rowCount);
}
}
connection.commit();
}
// The ExecuteSql request should only be sent once to Cloud Spanner.
assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
ExecuteSqlRequest executeRequest = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0);
assertEquals(QueryMode.NORMAL, executeRequest.getQueryMode());
assertEquals(SELECT_RANDOM.getSql(), executeRequest.getSql());
// 5. COMMIT
if (pgServer != null) {
List<DescribeMessage> describeMessages = getWireMessagesOfType(DescribeMessage.class);
assertEquals(1, describeMessages.size());
DescribeMessage describeMessage = describeMessages.get(0);
assertEquals(PreparedType.Portal, describeMessage.getType());
List<ExecuteMessage> executeMessages = getWireMessagesOfType(ExecuteMessage.class);
int expectedExecuteMessageCount = RANDOM_RESULTS_ROW_COUNT / fetchSize + ((RANDOM_RESULTS_ROW_COUNT % fetchSize) > 0 ? 1 : 0) + 2;
assertEquals(expectedExecuteMessageCount, executeMessages.size());
assertEquals("", executeMessages.get(0).getName());
for (ExecuteMessage executeMessage : executeMessages.subList(1, executeMessages.size() - 1)) {
assertEquals(describeMessage.getName(), executeMessage.getName());
assertEquals(fetchSize, executeMessage.getMaxRows());
}
assertEquals("", executeMessages.get(executeMessages.size() - 1).getName());
List<ParseMessage> parseMessages = getWireMessagesOfType(ParseMessage.class);
assertEquals(3, parseMessages.size());
assertEquals("BEGIN", parseMessages.get(0).getStatement().getSql());
assertEquals(SELECT_RANDOM.getSql(), parseMessages.get(1).getStatement().getSql());
assertEquals("COMMIT", parseMessages.get(2).getStatement().getSql());
}
}
use of com.google.cloud.spanner.pgadapter.wireprotocol.ExecuteMessage in project pgadapter by GoogleCloudPlatform.
the class JdbcMockServerTest method testCursorFailsHalfway.
@Test
public void testCursorFailsHalfway() throws SQLException {
mockSpanner.setExecuteStreamingSqlExecutionTime(SimulatedExecutionTime.ofStreamException(Status.DATA_LOSS.asRuntimeException(), 2));
try (Connection connection = DriverManager.getConnection(createUrl())) {
connection.setAutoCommit(false);
try (PreparedStatement statement = connection.prepareStatement(SELECT_FIVE_ROWS.getSql())) {
// Fetch one row at a time from the PG server.
statement.setFetchSize(1);
try (ResultSet resultSet = statement.executeQuery()) {
// The first row should succeed.
assertTrue(resultSet.next());
// The second row should fail.
assertThrows(SQLException.class, resultSet::next);
}
}
connection.rollback();
}
// The ExecuteSql request should only be sent once to Cloud Spanner.
assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
ExecuteSqlRequest executeRequest = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0);
assertEquals(QueryMode.NORMAL, executeRequest.getQueryMode());
assertEquals(SELECT_FIVE_ROWS.getSql(), executeRequest.getSql());
// The JDBC driver does not send a ROLLBACK
if (pgServer != null) {
List<DescribeMessage> describeMessages = getWireMessagesOfType(DescribeMessage.class);
assertEquals(1, describeMessages.size());
DescribeMessage describeMessage = describeMessages.get(0);
assertEquals(PreparedType.Portal, describeMessage.getType());
List<ExecuteMessage> executeMessages = getWireMessagesOfType(ExecuteMessage.class);
assertEquals(4, executeMessages.size());
assertEquals("", executeMessages.get(0).getName());
for (ExecuteMessage executeMessage : executeMessages.subList(1, executeMessages.size() - 1)) {
assertEquals(describeMessage.getName(), executeMessage.getName());
assertEquals(1, executeMessage.getMaxRows());
}
assertEquals("", executeMessages.get(executeMessages.size() - 1).getName());
List<ParseMessage> parseMessages = getWireMessagesOfType(ParseMessage.class);
assertEquals(3, parseMessages.size());
assertEquals("BEGIN", parseMessages.get(0).getStatement().getSql());
assertEquals(SELECT_FIVE_ROWS.getSql(), parseMessages.get(1).getStatement().getSql());
assertEquals("ROLLBACK", parseMessages.get(2).getStatement().getSql());
}
}
Aggregations