Search in sources :

Example 11 with MutationWriter

use of com.google.cloud.spanner.pgadapter.utils.MutationWriter in project pgadapter by GoogleCloudPlatform.

the class ProtocolTest method testCopyDataMessage.

@Test
public void testCopyDataMessage() throws Exception {
    byte[] messageMetadata = { 'd' };
    byte[] payload = "This is the payload".getBytes();
    byte[] length = intToBytes(4 + payload.length);
    byte[] value = Bytes.concat(messageMetadata, length, payload);
    DataInputStream inputStream = new DataInputStream(new ByteArrayInputStream(value));
    CopyStatement copyStatement = mock(CopyStatement.class);
    when(connectionHandler.getActiveStatement()).thenReturn(copyStatement);
    when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
    when(connectionHandler.getStatus()).thenReturn(ConnectionStatus.COPY_IN);
    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
    MutationWriter mw = mock(MutationWriter.class);
    when(copyStatement.getMutationWriter()).thenReturn(mw);
    WireMessage message = ControlMessage.create(connectionHandler);
    assertEquals(CopyDataMessage.class, message.getClass());
    assertArrayEquals(payload, ((CopyDataMessage) message).getPayload());
    CopyDataMessage messageSpy = (CopyDataMessage) spy(message);
    messageSpy.send();
    verify(mw).addCopyData(payload);
}
Also used : CopyDataMessage(com.google.cloud.spanner.pgadapter.wireprotocol.CopyDataMessage) ByteArrayInputStream(java.io.ByteArrayInputStream) CopyStatement(com.google.cloud.spanner.pgadapter.statements.CopyStatement) MutationWriter(com.google.cloud.spanner.pgadapter.utils.MutationWriter) WireMessage(com.google.cloud.spanner.pgadapter.wireprotocol.WireMessage) DataInputStream(java.io.DataInputStream) Test(org.junit.Test)

Example 12 with MutationWriter

use of com.google.cloud.spanner.pgadapter.utils.MutationWriter in project pgadapter by GoogleCloudPlatform.

the class ProtocolTest method testCopyDataRowLengthMismatchLimit.

@Test
public void testCopyDataRowLengthMismatchLimit() throws Exception {
    setupQueryInformationSchemaResults();
    byte[] payload = "1\t'one'\n2".getBytes();
    CopyStatement copyStatement = new CopyStatement(connectionHandler, mock(OptionsMetadata.class), parse("COPY keyvalue FROM STDIN;"));
    assertFalse(copyStatement.isExecuted());
    copyStatement.execute();
    assertTrue(copyStatement.isExecuted());
    MutationWriter mw = copyStatement.getMutationWriter();
    mw.addCopyData(payload);
    mw.close();
    SpannerException thrown = assertThrows(SpannerException.class, copyStatement::getUpdateCount);
    assertEquals(ErrorCode.INVALID_ARGUMENT, thrown.getErrorCode());
    assertEquals("INVALID_ARGUMENT: Invalid COPY data: Row length mismatched. Expected 2 columns, but only found 1", thrown.getMessage());
    copyStatement.close();
    File outputFile = new File("output.txt");
    if (outputFile.exists()) {
        assertTrue(outputFile.delete());
    }
}
Also used : CopyStatement(com.google.cloud.spanner.pgadapter.statements.CopyStatement) MutationWriter(com.google.cloud.spanner.pgadapter.utils.MutationWriter) OptionsMetadata(com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata) SpannerException(com.google.cloud.spanner.SpannerException) File(java.io.File) Test(org.junit.Test)

Example 13 with MutationWriter

use of com.google.cloud.spanner.pgadapter.utils.MutationWriter in project pgadapter by GoogleCloudPlatform.

the class ProtocolTest method testCopyDoneMessage.

@Test
public void testCopyDoneMessage() throws Exception {
    byte[] messageMetadata = { 'c' };
    byte[] length = intToBytes(4);
    byte[] value = Bytes.concat(messageMetadata, length);
    DataInputStream inputStream = new DataInputStream(new ByteArrayInputStream(value));
    ByteArrayOutputStream result = new ByteArrayOutputStream();
    DataOutputStream outputStream = new DataOutputStream(result);
    CopyStatement copyStatement = mock(CopyStatement.class);
    when(connectionHandler.getActiveStatement()).thenReturn(copyStatement);
    when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
    when(connectionHandler.getStatus()).thenReturn(ConnectionStatus.COPY_IN);
    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
    MutationWriter mb = mock(MutationWriter.class);
    when(copyStatement.getMutationWriter()).thenReturn(mb);
    WireMessage message = ControlMessage.create(connectionHandler);
    assertEquals(CopyDoneMessage.class, message.getClass());
    CopyDoneMessage messageSpy = (CopyDoneMessage) spy(message);
    doReturn(false).when(messageSpy).sendSpannerResult(anyInt(), any(IntermediateStatement.class), any(QueryMode.class), anyLong());
    messageSpy.send();
    verify(messageSpy).sendSpannerResult(0, copyStatement, QueryMode.SIMPLE, 0L);
}
Also used : ByteArrayInputStream(java.io.ByteArrayInputStream) CopyStatement(com.google.cloud.spanner.pgadapter.statements.CopyStatement) DataOutputStream(java.io.DataOutputStream) QueryMode(com.google.cloud.spanner.pgadapter.ConnectionHandler.QueryMode) MutationWriter(com.google.cloud.spanner.pgadapter.utils.MutationWriter) WireMessage(com.google.cloud.spanner.pgadapter.wireprotocol.WireMessage) IntermediateStatement(com.google.cloud.spanner.pgadapter.statements.IntermediateStatement) CopyDoneMessage(com.google.cloud.spanner.pgadapter.wireprotocol.CopyDoneMessage) ByteArrayOutputStream(java.io.ByteArrayOutputStream) DataInputStream(java.io.DataInputStream) Test(org.junit.Test)

Aggregations

MutationWriter (com.google.cloud.spanner.pgadapter.utils.MutationWriter)13 CopyStatement (com.google.cloud.spanner.pgadapter.statements.CopyStatement)10 Test (org.junit.Test)10 SpannerException (com.google.cloud.spanner.SpannerException)5 OptionsMetadata (com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata)5 WireMessage (com.google.cloud.spanner.pgadapter.wireprotocol.WireMessage)3 ByteArrayInputStream (java.io.ByteArrayInputStream)3 DataInputStream (java.io.DataInputStream)3 File (java.io.File)3 IntermediateStatement (com.google.cloud.spanner.pgadapter.statements.IntermediateStatement)2 ReadyResponse (com.google.cloud.spanner.pgadapter.wireoutput.ReadyResponse)2 ByteArrayOutputStream (java.io.ByteArrayOutputStream)2 DataOutputStream (java.io.DataOutputStream)2 DatabaseClient (com.google.cloud.spanner.DatabaseClient)1 ResultSet (com.google.cloud.spanner.ResultSet)1 Statement (com.google.cloud.spanner.Statement)1 QueryMode (com.google.cloud.spanner.pgadapter.ConnectionHandler.QueryMode)1 IntermediatePortalStatement (com.google.cloud.spanner.pgadapter.statements.IntermediatePortalStatement)1 IntermediatePreparedStatement (com.google.cloud.spanner.pgadapter.statements.IntermediatePreparedStatement)1 ErrorResponse (com.google.cloud.spanner.pgadapter.wireoutput.ErrorResponse)1