Examples with WireMessage com.google.cloud.spanner.pgadapter.wireprotocol.WireMessage used on opensource projects

Example 11 with WireMessage

use of com.google.cloud.spanner.pgadapter.wireprotocol.WireMessage in project pgadapter by GoogleCloudPlatform.

the class ProtocolTest method testParseMessageAcceptsUntypedParameter.

@Test
public void testParseMessageAcceptsUntypedParameter() throws Exception {
    byte[] messageMetadata = { 'P' };
    String statementName = "some statement\0";
    String payload = "SELECT * FROM users WHERE name = $1 /*This is a comment*/ --this is another comment\0";
    byte[] parameterCount = { 0, 1 };
    // Unspecifed parameter type.
    byte[] parameters = intToBytes(0);
    byte[] length = intToBytes(4 + statementName.length() + payload.length() + parameterCount.length + parameters.length);
    byte[] value = Bytes.concat(messageMetadata, length, statementName.getBytes(), payload.getBytes(), parameterCount, parameters);
    int[] expectedParameterDataTypes = new int[] { 0 };
    String expectedSQL = "SELECT * FROM users WHERE name = $1";
    String expectedMessageName = "some statement";
    DataInputStream inputStream = new DataInputStream(new ByteArrayInputStream(value));
    ByteArrayOutputStream result = new ByteArrayOutputStream();
    DataOutputStream outputStream = new DataOutputStream(result);
    when(connectionHandler.getServer()).thenReturn(server);
    when(server.getOptions()).thenReturn(options);
    when(connectionHandler.getSpannerConnection()).thenReturn(connection);
    when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
    WireMessage message = ControlMessage.create(connectionHandler);
    assertEquals(ParseMessage.class, message.getClass());
    assertEquals(expectedMessageName, ((ParseMessage) message).getName());
    assertEquals(expectedSQL, ((ParseMessage) message).getStatement().getSql());
    assertArrayEquals(expectedParameterDataTypes, ((ParseMessage) message).getStatement().getParameterDataTypes());
    when(connectionHandler.hasStatement(anyString())).thenReturn(false);
    message.send();
    verify(connectionHandler).registerStatement(expectedMessageName, ((ParseMessage) message).getStatement());
    // ParseCompleteResponse
    DataInputStream outputResult = inputStreamFromOutputStream(result);
    assertEquals('1', outputResult.readByte());
    assertEquals(4, outputResult.readInt());
}

Example 12 with WireMessage

use of com.google.cloud.spanner.pgadapter.wireprotocol.WireMessage 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);
}

Example 13 with WireMessage

use of com.google.cloud.spanner.pgadapter.wireprotocol.WireMessage in project pgadapter by GoogleCloudPlatform.

the class ProtocolTest method testMultipleCopyDataMessages.

@Test
public void testMultipleCopyDataMessages() throws Exception {
    when(connectionHandler.getSpannerConnection()).thenReturn(connection);
    when(connectionHandler.getStatus()).thenReturn(ConnectionStatus.COPY_IN);
    byte[] messageMetadata = { 'd' };
    byte[] payload1 = "1\t'one'\n2\t".getBytes();
    byte[] payload2 = "'two'\n3\t'th".getBytes();
    byte[] payload3 = "ree'\n4\t'four'\n".getBytes();
    byte[] length1 = intToBytes(4 + payload1.length);
    byte[] length2 = intToBytes(4 + payload2.length);
    byte[] length3 = intToBytes(4 + payload3.length);
    byte[] value1 = Bytes.concat(messageMetadata, length1, payload1);
    byte[] value2 = Bytes.concat(messageMetadata, length2, payload2);
    byte[] value3 = Bytes.concat(messageMetadata, length3, payload3);
    DataInputStream inputStream1 = new DataInputStream(new ByteArrayInputStream(value1));
    DataInputStream inputStream2 = new DataInputStream(new ByteArrayInputStream(value2));
    DataInputStream inputStream3 = new DataInputStream(new ByteArrayInputStream(value3));
    ResultSet spannerType = mock(ResultSet.class);
    when(spannerType.getString("column_name")).thenReturn("key", "value");
    when(spannerType.getString("data_type")).thenReturn("bigint", "character varying");
    when(spannerType.next()).thenReturn(true, true, false);
    when(connection.executeQuery(any(Statement.class))).thenReturn(spannerType);
    CopyStatement copyStatement = new CopyStatement(connectionHandler, options, parse("COPY keyvalue FROM STDIN;"));
    copyStatement.execute();
    when(connectionHandler.getActiveStatement()).thenReturn(copyStatement);
    when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
    {
        when(connectionMetadata.getInputStream()).thenReturn(inputStream1);
        WireMessage message = ControlMessage.create(connectionHandler);
        assertEquals(CopyDataMessage.class, message.getClass());
        assertArrayEquals(payload1, ((CopyDataMessage) message).getPayload());
        CopyDataMessage copyDataMessage = (CopyDataMessage) message;
        copyDataMessage.send();
    }
    {
        when(connectionMetadata.getInputStream()).thenReturn(inputStream2);
        WireMessage message = ControlMessage.create(connectionHandler);
        assertEquals(CopyDataMessage.class, message.getClass());
        assertArrayEquals(payload2, ((CopyDataMessage) message).getPayload());
        CopyDataMessage copyDataMessage = (CopyDataMessage) message;
        copyDataMessage.send();
    }
    {
        when(connectionMetadata.getInputStream()).thenReturn(inputStream3);
        WireMessage message = ControlMessage.create(connectionHandler);
        assertEquals(CopyDataMessage.class, message.getClass());
        assertArrayEquals(payload3, ((CopyDataMessage) message).getPayload());
        CopyDataMessage copyDataMessage = (CopyDataMessage) message;
        copyDataMessage.send();
    }
}

Example 14 with WireMessage

use of com.google.cloud.spanner.pgadapter.wireprotocol.WireMessage in project pgadapter by GoogleCloudPlatform.

the class ProtocolTest method testCloseStatementMessage.

@Test
public void testCloseStatementMessage() throws Exception {
    byte[] messageMetadata = { 'C' };
    byte[] statementType = { 'S' };
    String statementName = "some statement\0";
    byte[] length = intToBytes(4 + statementType.length + statementName.length());
    byte[] value = Bytes.concat(messageMetadata, length, statementType, statementName.getBytes());
    String expectedStatementName = "some statement";
    PreparedType expectedType = PreparedType.Statement;
    DataInputStream inputStream = new DataInputStream(new ByteArrayInputStream(value));
    ByteArrayOutputStream result = new ByteArrayOutputStream();
    DataOutputStream outputStream = new DataOutputStream(result);
    when(connectionHandler.getStatement(anyString())).thenReturn(intermediatePortalStatement);
    when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
    WireMessage message = ControlMessage.create(connectionHandler);
    assertEquals(CloseMessage.class, message.getClass());
    assertEquals(expectedStatementName, ((CloseMessage) message).getName());
    assertEquals(expectedType, ((CloseMessage) message).getType());
    verify(connectionHandler).getStatement("some statement");
    message.send();
    verify(connectionHandler).closeStatement(expectedStatementName);
    // CloseResponse
    DataInputStream outputResult = inputStreamFromOutputStream(result);
    assertEquals('3', outputResult.readByte());
    assertEquals(4, outputResult.readInt());
}

Example 15 with WireMessage

use of com.google.cloud.spanner.pgadapter.wireprotocol.WireMessage in project pgadapter by GoogleCloudPlatform.

the class ProtocolTest method testParseMessageWorksIfNameIsEmpty.

@Test
public void testParseMessageWorksIfNameIsEmpty() throws Exception {
    byte[] messageMetadata = { 'P' };
    String statementName = "\0";
    String payload = "SELECT * FROM users WHERE name = $1 /*This is a comment*/ --this is another comment\0";
    byte[] parameterCount = { 0, 1 };
    byte[] parameters = intToBytes(1002);
    byte[] length = intToBytes(4 + statementName.length() + payload.length() + parameterCount.length + parameters.length);
    byte[] value = Bytes.concat(messageMetadata, length, statementName.getBytes(), payload.getBytes(), parameterCount, parameters);
    ByteArrayOutputStream result = new ByteArrayOutputStream();
    DataOutputStream outputStream = new DataOutputStream(result);
    DataInputStream inputStream = new DataInputStream(new ByteArrayInputStream(value));
    when(connectionHandler.getServer()).thenReturn(server);
    when(server.getOptions()).thenReturn(options);
    when(connectionHandler.getSpannerConnection()).thenReturn(connection);
    when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
    WireMessage message = ControlMessage.create(connectionHandler);
    message.send();
}