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);
}
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());
}
}
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);
}
Aggregations