Search in sources :

Example 1 with ExecuteMessage

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);
}
Also used : ByteArrayInputStream(java.io.ByteArrayInputStream) ExecuteMessage(com.google.cloud.spanner.pgadapter.wireprotocol.ExecuteMessage) DataOutputStream(java.io.DataOutputStream) QueryMode(com.google.cloud.spanner.pgadapter.ConnectionHandler.QueryMode) WireMessage(com.google.cloud.spanner.pgadapter.wireprotocol.WireMessage) ArgumentMatchers.anyString(org.mockito.ArgumentMatchers.anyString) ByteArrayOutputStream(java.io.ByteArrayOutputStream) DataInputStream(java.io.DataInputStream) IntermediatePortalStatement(com.google.cloud.spanner.pgadapter.statements.IntermediatePortalStatement) Test(org.junit.Test)

Example 2 with ExecuteMessage

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());
    }
}
Also used : ExecuteSqlRequest(com.google.spanner.v1.ExecuteSqlRequest) ParseMessage(com.google.cloud.spanner.pgadapter.wireprotocol.ParseMessage) ExecuteMessage(com.google.cloud.spanner.pgadapter.wireprotocol.ExecuteMessage) DescribeMessage(com.google.cloud.spanner.pgadapter.wireprotocol.DescribeMessage) Connection(java.sql.Connection) ResultSet(java.sql.ResultSet) PreparedStatement(java.sql.PreparedStatement) Test(org.junit.Test)

Example 3 with ExecuteMessage

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());
    }
}
Also used : ExecuteSqlRequest(com.google.spanner.v1.ExecuteSqlRequest) ParseMessage(com.google.cloud.spanner.pgadapter.wireprotocol.ParseMessage) ExecuteMessage(com.google.cloud.spanner.pgadapter.wireprotocol.ExecuteMessage) DescribeMessage(com.google.cloud.spanner.pgadapter.wireprotocol.DescribeMessage) Connection(java.sql.Connection) ResultSet(java.sql.ResultSet) PreparedStatement(java.sql.PreparedStatement) Test(org.junit.Test)

Example 4 with ExecuteMessage

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());
    }
}
Also used : ExecuteSqlRequest(com.google.spanner.v1.ExecuteSqlRequest) ParseMessage(com.google.cloud.spanner.pgadapter.wireprotocol.ParseMessage) ExecuteMessage(com.google.cloud.spanner.pgadapter.wireprotocol.ExecuteMessage) DescribeMessage(com.google.cloud.spanner.pgadapter.wireprotocol.DescribeMessage) Connection(java.sql.Connection) ResultSet(java.sql.ResultSet) PreparedStatement(java.sql.PreparedStatement) Test(org.junit.Test)

Aggregations

ExecuteMessage (com.google.cloud.spanner.pgadapter.wireprotocol.ExecuteMessage)4 Test (org.junit.Test)4 DescribeMessage (com.google.cloud.spanner.pgadapter.wireprotocol.DescribeMessage)3 ParseMessage (com.google.cloud.spanner.pgadapter.wireprotocol.ParseMessage)3 ExecuteSqlRequest (com.google.spanner.v1.ExecuteSqlRequest)3 Connection (java.sql.Connection)3 PreparedStatement (java.sql.PreparedStatement)3 ResultSet (java.sql.ResultSet)3 QueryMode (com.google.cloud.spanner.pgadapter.ConnectionHandler.QueryMode)1 IntermediatePortalStatement (com.google.cloud.spanner.pgadapter.statements.IntermediatePortalStatement)1 WireMessage (com.google.cloud.spanner.pgadapter.wireprotocol.WireMessage)1 ByteArrayInputStream (java.io.ByteArrayInputStream)1 ByteArrayOutputStream (java.io.ByteArrayOutputStream)1 DataInputStream (java.io.DataInputStream)1 DataOutputStream (java.io.DataOutputStream)1 ArgumentMatchers.anyString (org.mockito.ArgumentMatchers.anyString)1