Examples with MutationWriter com.google.cloud.spanner.pgadapter.utils.MutationWriter used on opensource projects

Example 1 with MutationWriter

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

Example 2 with MutationWriter

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

the class CopyStatement method execute.

@Override
public void execute() {
    this.executedCount++;
    try {
        parseCopyStatement();
        queryInformationSchema();
        setParserFormat(this.options);
        mutationWriter = new MutationWriter(getTransactionMode(), connection, options.getTableName(), getTableColumns(), indexedColumnsCount, getParserFormat(), hasHeader());
        updateCount = executor.submit(mutationWriter);
    } catch (Exception e) {
        SpannerException spannerException = SpannerExceptionFactory.asSpannerException(e);
        handleExecutionException(spannerException);
    }
}

Example 3 with MutationWriter

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

the class CopyDoneMessage method sendPayload.

@Override
protected void sendPayload() throws Exception {
    // If backend error occurred during copy-in mode, drop any subsequent CopyDone messages.
    if (this.statement != null) {
        MutationWriter mutationWriter = this.statement.getMutationWriter();
        statement.close();
        if (!statement.hasException(0)) {
            try {
                long rowCount = this.statement.getUpdateCount();
                // Set the row count of number of rows copied.
                statement.setUpdateCount(0, rowCount);
                this.sendSpannerResult(0, this.statement, QueryMode.SIMPLE, 0L);
                this.outputStream.flush();
            } catch (Exception e) {
                // Spanner returned an error when trying to commit the batch of mutations.
                mutationWriter.writeErrorFile(e);
                mutationWriter.closeErrorFile();
                this.connection.setStatus(ConnectionStatus.IDLE);
                this.connection.removeActiveStatement(this.statement);
                throw e;
            }
        } else {
            mutationWriter.closeErrorFile();
        }
    }
    this.connection.setStatus(ConnectionStatus.IDLE);
    this.connection.removeActiveStatement(this.statement);
    new ReadyResponse(this.outputStream, ReadyResponse.Status.IDLE).send();
}

Example 4 with MutationWriter

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

the class CopyFailMessage method sendPayload.

@Override
protected void sendPayload() throws Exception {
    // If backend error occurred during copy-in mode, drop any subsequent CopyFail messages.
    if (this.statement != null) {
        MutationWriter mutationWriter = this.statement.getMutationWriter();
        mutationWriter.rollback();
        mutationWriter.closeErrorFile();
        statement.close();
        if (!statement.hasException(0)) {
            new ErrorResponse(this.outputStream, new Exception(this.errorMessage), State.IOError).send();
        }
    }
    this.connection.setStatus(ConnectionStatus.IDLE);
    this.connection.removeActiveStatement(this.statement);
    new ReadyResponse(this.outputStream, ReadyResponse.Status.IDLE).send();
}

Example 5 with MutationWriter

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