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

Example 11 with CopyStatement

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

Example 12 with CopyStatement

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

the class QueryMessage method handleQuery.

/**
 * Simple Query handler, which examined the state of the statement and processes accordingly (if
 * error, handle error, otherwise sends the result and if contains result set, send row
 * description)
 *
 * @throws Exception if handling the query fails
 */
public void handleQuery() throws Exception {
    // Skip unexecuted statements, as no response needs be returned
    for (int index = 0; index < statement.getExecutedCount(); index++) {
        if (this.statement.hasException(index)) {
            new ErrorResponse(this.outputStream, this.statement.getException(index), State.InternalError).send();
        } else if (this.statement.getCommand(index).equalsIgnoreCase(COPY)) {
            CopyStatement copyStatement = (CopyStatement) this.statement;
            new CopyInResponse(this.outputStream, copyStatement.getTableColumns().size(), copyStatement.getFormatCode()).send();
            this.connection.setStatus(ConnectionStatus.COPY_IN);
            // Return early as we do not respond with CommandComplete after a COPY command.
            return;
        } else if ("".equals(this.statement.getCommandTag(index))) {
            new EmptyQueryResponse(outputStream).send();
        } else {
            if (this.statement.containsResultSet(index)) {
                new RowDescriptionResponse(this.outputStream, this.statement, this.statement.getStatementResult(index), this.connection.getServer().getOptions(), QueryMode.SIMPLE).send(false);
            }
            this.sendSpannerResult(index, this.statement, QueryMode.SIMPLE, 0L);
        }
    }
    if (connection.getSpannerConnection().isInTransaction() && connection.getStatus() == ConnectionStatus.TRANSACTION_ABORTED) {
        // Actively rollback the aborted transaction but still block clients
        // Clear any statement tags, as these are not allowed for rollbacks.
        connection.getSpannerConnection().setStatementTag(null);
        connection.getSpannerConnection().rollback();
    }
    new ReadyResponse(this.outputStream, connection.getStatus().getReadyResponseStatus()).send();
    this.connection.cleanUp(this.statement);
}