use of com.google.cloud.spanner.pgadapter.statements.CopyStatement in project pgadapter by GoogleCloudPlatform.
the class StatementTest method testCopyBuildMutation.
@Test
public void testCopyBuildMutation() throws Exception {
Mockito.when(connectionHandler.getSpannerConnection()).thenReturn(connection);
Mockito.when(statementResult.getResultType()).thenReturn(StatementResult.ResultType.UPDATE_COUNT);
Mockito.when(statementResult.getUpdateCount()).thenReturn(1L);
ResultSet spannerType = Mockito.mock(ResultSet.class);
Mockito.when(spannerType.getString("column_name")).thenReturn("key", "value");
Mockito.when(spannerType.getString("data_type")).thenReturn("bigint", "character varying");
Mockito.when(spannerType.next()).thenReturn(true, true, false);
Mockito.when(connection.executeQuery(any(Statement.class))).thenReturn(spannerType);
CopyStatement statement = new CopyStatement("COPY keyvalue FROM STDIN;", connection);
statement.execute();
byte[] payload = "2\t3\n".getBytes();
MutationWriter mw = statement.getMutationWriter();
mw.addCopyData(connectionHandler, payload);
mw.buildMutationList(connectionHandler);
Assert.assertEquals(statement.getFormatType(), "TEXT");
Assert.assertEquals(statement.getDelimiterChar(), '\t');
Assert.assertEquals(statement.getMutationWriter().getMutations().toString(), "[insert(keyvalue{key=2,value=3})]");
statement.close();
Mockito.verify(resultSet, Mockito.times(0)).close();
}
use of com.google.cloud.spanner.pgadapter.statements.CopyStatement in project pgadapter by GoogleCloudPlatform.
the class ProtocolTest method testCopyFailMessage.
@Test
public void testCopyFailMessage() throws Exception {
byte[] messageMetadata = { 'f' };
byte[] errorMessage = "Error Message\0".getBytes();
byte[] length = intToBytes(4 + errorMessage.length);
byte[] value = Bytes.concat(messageMetadata, length, errorMessage);
DataInputStream inputStream = new DataInputStream(new ByteArrayInputStream(value));
ByteArrayOutputStream result = new ByteArrayOutputStream();
DataOutputStream outputStream = new DataOutputStream(result);
String expectedErrorMessage = "Error Message";
CopyStatement copyStatement = mock(CopyStatement.class);
MutationWriter mutationWriter = mock(MutationWriter.class);
when(copyStatement.getMutationWriter()).thenReturn(mutationWriter);
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);
WireMessage message = ControlMessage.create(connectionHandler);
assertEquals(CopyFailMessage.class, message.getClass());
assertEquals(expectedErrorMessage, ((CopyFailMessage) message).getErrorMessage());
message.send();
verify(mutationWriter).rollback();
}
use of com.google.cloud.spanner.pgadapter.statements.CopyStatement in project pgadapter by GoogleCloudPlatform.
the class ProtocolTest method testCopyBatchSizeLimit.
@Test
public void testCopyBatchSizeLimit() throws Exception {
setupQueryInformationSchemaResults();
byte[] payload = Files.readAllBytes(Paths.get("./src/test/resources/batch-size-test.txt"));
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();
// Inject a mock DatabaseClient for now.
// TODO: Fix this once we can use the Connection API.
Field databaseClientField = MutationWriter.class.getDeclaredField("databaseClient");
databaseClientField.setAccessible(true);
databaseClientField.set(mw, mock(DatabaseClient.class));
mw.addCopyData(payload);
mw.close();
assertEquals("TEXT", copyStatement.getFormatType());
assertEquals('\t', copyStatement.getDelimiterChar());
assertFalse(copyStatement.hasException(0));
assertEquals(12L, copyStatement.getUpdateCount().longValue());
assertEquals(12L, mw.getRowCount());
copyStatement.close();
}
use of com.google.cloud.spanner.pgadapter.statements.CopyStatement in project pgadapter by GoogleCloudPlatform.
the class ProtocolTest method testCopyResumeErrorStartOutputFile.
@Test
public void testCopyResumeErrorStartOutputFile() throws Exception {
setupQueryInformationSchemaResults();
byte[] payload = Files.readAllBytes(Paths.get("./src/test/resources/test-copy-start-output.txt"));
// Pre-emptively try to delete the output file if it is lingering from a previous test run.
// TODO: Make the output file for COPY configurable, so we can use a temp file for this.
File outputFile = new File("output.txt");
if (outputFile.exists()) {
assertTrue(outputFile.delete());
}
CopyStatement copyStatement = new CopyStatement(connectionHandler, options, parse("COPY keyvalue FROM STDIN;"));
assertFalse(copyStatement.isExecuted());
copyStatement.execute();
assertTrue(copyStatement.isExecuted());
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:\"'1'\"", thrown.getMessage());
outputFile = new File("output.txt");
assertTrue(outputFile.exists());
assertTrue(outputFile.isFile());
assertTrue(outputFile.delete());
copyStatement.close();
}
use of com.google.cloud.spanner.pgadapter.statements.CopyStatement 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();
}
}
Aggregations