Examples with CopyStatement com.google.cloud.spanner.pgadapter.statements.CopyStatement used on opensource projects

Example 6 with CopyStatement

use of com.google.cloud.spanner.pgadapter.statements.CopyStatement in project pgadapter by GoogleCloudPlatform.

the class ProtocolTest method testCopyResumeErrorOutputFile.

@Test
public void testCopyResumeErrorOutputFile() throws Exception {
    setupQueryInformationSchemaResults();
    byte[] payload = Files.readAllBytes(Paths.get("./src/test/resources/test-copy-output.txt"));
    CopyStatement copyStatement = new CopyStatement(connectionHandler, mock(OptionsMetadata.class), parse("COPY keyvalue FROM STDIN;"));
    assertFalse(copyStatement.isExecuted());
    copyStatement.execute();
    assertTrue(copyStatement.isExecuted());
    File outputFile = new File("output.txt");
    if (outputFile.exists()) {
        assertTrue(outputFile.delete());
    }
    MutationWriter mw = copyStatement.getMutationWriter();
    mw.addCopyData(payload);
    SpannerException thrown = assertThrows(SpannerException.class, copyStatement::getUpdateCount);
    assertEquals(ErrorCode.INVALID_ARGUMENT, thrown.getErrorCode());
    assertEquals("INVALID_ARGUMENT: Invalid input syntax for type INT64:\"'5'\"", thrown.getMessage());
    outputFile = new File("output.txt");
    assertTrue(outputFile.exists());
    assertTrue(outputFile.isFile());
    assertTrue(outputFile.delete());
    copyStatement.close();
}

Example 7 with CopyStatement

use of com.google.cloud.spanner.pgadapter.statements.CopyStatement in project pgadapter by GoogleCloudPlatform.

the class ProtocolTest method testCopyFromFilePipe.

@Test
public void testCopyFromFilePipe() throws Exception {
    setupQueryInformationSchemaResults();
    byte[] payload = Files.readAllBytes(Paths.get("./src/test/resources/small-file-test.txt"));
    CopyStatement copyStatement = new CopyStatement(connectionHandler, mock(OptionsMetadata.class), parse("COPY keyvalue FROM STDIN;"));
    copyStatement.execute();
    MutationWriter mw = copyStatement.getMutationWriter();
    mw.addCopyData(payload);
    assertEquals("TEXT", copyStatement.getFormatType());
    assertEquals('\t', copyStatement.getDelimiterChar());
    copyStatement.close();
    verify(resultSet, never()).close();
}

Example 8 with CopyStatement

use of com.google.cloud.spanner.pgadapter.statements.CopyStatement in project pgadapter by GoogleCloudPlatform.

the class StatementTest method testCopyInvalidBuildMutation.

@Test
public void testCopyInvalidBuildMutation() throws Exception {
    when(connectionHandler.getSpannerConnection()).thenReturn(connection);
    setupQueryInformationSchemaResults();
    CopyStatement statement = new CopyStatement(connectionHandler, mock(OptionsMetadata.class), parse("COPY keyvalue FROM STDIN;"));
    statement.execute();
    byte[] payload = "2 3\n".getBytes();
    MutationWriter mutationWriter = statement.getMutationWriter();
    mutationWriter.addCopyData(payload);
    SpannerException thrown = assertThrows(SpannerException.class, statement::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());
    statement.close();
    verify(resultSet, never()).close();
}

Example 9 with CopyStatement

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

Example 10 with CopyStatement

use of com.google.cloud.spanner.pgadapter.statements.CopyStatement 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());
    }
}