Examples with Connection com.google.cloud.spanner.connection.Connection used on opensource projects

Example 11 with Connection

use of com.google.cloud.spanner.connection.Connection in project pgadapter by GoogleCloudPlatform.

the class ProtocolTest method setupQueryInformationSchemaResults.

private void setupQueryInformationSchemaResults() {
    when(connectionHandler.getSpannerConnection()).thenReturn(connection);
    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(ArgumentMatchers.argThat(statement -> statement != null && statement.getSql().startsWith("SELECT column_name")))).thenReturn(spannerType);
    ResultSet countResult = mock(ResultSet.class);
    when(countResult.getLong(0)).thenReturn(2L);
    when(countResult.next()).thenReturn(true, false);
    when(connection.executeQuery(ArgumentMatchers.argThat(statement -> statement != null && statement.getSql().startsWith("SELECT COUNT(*)")))).thenReturn(countResult);
}

Example 12 with Connection

use of com.google.cloud.spanner.connection.Connection in project pgadapter by GoogleCloudPlatform.

the class ConnectionHandler method connectToSpanner.

@InternalApi
public void connectToSpanner(String database) {
    OptionsMetadata options = getServer().getOptions();
    String uri = options.hasDefaultConnectionUrl() ? options.getDefaultConnectionUrl() : options.buildConnectionURL(database);
    if (uri.startsWith("jdbc:")) {
        uri = uri.substring("jdbc:".length());
    }
    uri = appendPropertiesToUrl(uri, getServer().getProperties());
    if (System.getProperty(CHANNEL_PROVIDER_PROPERTY) != null) {
        uri = uri + ";" + ConnectionOptions.CHANNEL_PROVIDER_PROPERTY_NAME + "=" + System.getProperty(CHANNEL_PROVIDER_PROPERTY);
        // This forces the connection to use NoCredentials.
        uri = uri + ";usePlainText=true";
        try {
            Class.forName(System.getProperty(CHANNEL_PROVIDER_PROPERTY));
        } catch (ClassNotFoundException e) {
            throw SpannerExceptionFactory.newSpannerException(ErrorCode.INVALID_ARGUMENT, "Unknown or invalid channel provider: " + System.getProperty(CHANNEL_PROVIDER_PROPERTY));
        }
    }
    ConnectionOptions connectionOptions = ConnectionOptions.newBuilder().setUri(uri).build();
    Connection spannerConnection = connectionOptions.getConnection();
    try {
        // invalid, for example as a result of the credentials being wrong.
        if (spannerConnection.getDialect() != Dialect.POSTGRESQL) {
            throw SpannerExceptionFactory.newSpannerException(ErrorCode.INVALID_ARGUMENT, String.format("The database uses dialect %s. Currently PGAdapter only supports connections to PostgreSQL dialect databases. " + "These can be created using https://cloud.google.com/spanner/docs/quickstart-console#postgresql", spannerConnection.getDialect()));
        }
    } catch (SpannerException e) {
        spannerConnection.close();
        throw e;
    }
    this.spannerConnection = spannerConnection;
}

Example 13 with Connection

use of com.google.cloud.spanner.connection.Connection in project java-spanner-jdbc by googleapis.

the class JdbcStatementTest method createStatement.

@SuppressWarnings("unchecked")
private JdbcStatement createStatement() throws SQLException {
    Connection spanner = mock(Connection.class);
    when(spanner.getDialect()).thenReturn(dialect);
    com.google.cloud.spanner.ResultSet resultSet = mock(com.google.cloud.spanner.ResultSet.class);
    when(resultSet.next()).thenReturn(true, false);
    when(resultSet.getColumnType(0)).thenReturn(Type.int64());
    when(resultSet.getLong(0)).thenReturn(1L);
    StatementResult selectResult = mock(StatementResult.class);
    when(selectResult.getResultType()).thenReturn(ResultType.RESULT_SET);
    when(selectResult.getResultSet()).thenReturn(resultSet);
    when(spanner.execute(com.google.cloud.spanner.Statement.of(SELECT))).thenReturn(selectResult);
    StatementResult updateResult = mock(StatementResult.class);
    when(updateResult.getResultType()).thenReturn(ResultType.UPDATE_COUNT);
    when(updateResult.getUpdateCount()).thenReturn(1L);
    when(spanner.execute(com.google.cloud.spanner.Statement.of(UPDATE))).thenReturn(updateResult);
    StatementResult largeUpdateResult = mock(StatementResult.class);
    when(largeUpdateResult.getResultType()).thenReturn(ResultType.UPDATE_COUNT);
    when(largeUpdateResult.getUpdateCount()).thenReturn(Integer.MAX_VALUE + 1L);
    when(spanner.execute(com.google.cloud.spanner.Statement.of(LARGE_UPDATE))).thenReturn(largeUpdateResult);
    StatementResult ddlResult = mock(StatementResult.class);
    when(ddlResult.getResultType()).thenReturn(ResultType.NO_RESULT);
    when(spanner.execute(com.google.cloud.spanner.Statement.of(DDL))).thenReturn(ddlResult);
    when(spanner.executeQuery(com.google.cloud.spanner.Statement.of(SELECT))).thenReturn(resultSet);
    when(spanner.executeQuery(com.google.cloud.spanner.Statement.of(UPDATE))).thenThrow(SpannerExceptionFactory.newSpannerException(ErrorCode.INVALID_ARGUMENT, "not a query"));
    when(spanner.executeQuery(com.google.cloud.spanner.Statement.of(DDL))).thenThrow(SpannerExceptionFactory.newSpannerException(ErrorCode.INVALID_ARGUMENT, "not a query"));
    when(spanner.executeUpdate(com.google.cloud.spanner.Statement.of(UPDATE))).thenReturn(1L);
    when(spanner.executeUpdate(com.google.cloud.spanner.Statement.of(SELECT))).thenThrow(SpannerExceptionFactory.newSpannerException(ErrorCode.INVALID_ARGUMENT, "not an update"));
    when(spanner.executeUpdate(com.google.cloud.spanner.Statement.of(DDL))).thenThrow(SpannerExceptionFactory.newSpannerException(ErrorCode.INVALID_ARGUMENT, "not an update"));
    when(spanner.executeBatchUpdate(anyList())).thenAnswer((Answer<long[]>) invocation -> {
        List<com.google.cloud.spanner.Statement> statements = (List<com.google.cloud.spanner.Statement>) invocation.getArguments()[0];
        if (statements.isEmpty() || AbstractStatementParser.getInstance(dialect).isDdlStatement(statements.get(0).getSql())) {
            return new long[0];
        }
        long[] res = new long[((List<com.google.cloud.spanner.Statement>) invocation.getArguments()[0]).size()];
        Arrays.fill(res, 1L);
        return res;
    });
    JdbcConnection connection = mock(JdbcConnection.class);
    when(connection.getDialect()).thenReturn(dialect);
    when(connection.getParser()).thenReturn(AbstractStatementParser.getInstance(dialect));
    when(connection.getSpannerConnection()).thenReturn(spanner);
    return new JdbcStatement(connection);
}

Example 14 with Connection

use of com.google.cloud.spanner.connection.Connection in project java-spanner-jdbc by googleapis.

the class JdbcStatementTest method testInternalExecuteUpdate.

@Test
public void testInternalExecuteUpdate() throws SQLException {
    JdbcConnection connection = mock(JdbcConnection.class);
    when(connection.getDialect()).thenReturn(dialect);
    Connection spannerConnection = mock(Connection.class);
    when(connection.getSpannerConnection()).thenReturn(spannerConnection);
    com.google.cloud.spanner.Statement updateStatement = com.google.cloud.spanner.Statement.of(UPDATE);
    com.google.cloud.spanner.Statement largeUpdateStatement = com.google.cloud.spanner.Statement.of(LARGE_UPDATE);
    when(spannerConnection.executeUpdate(updateStatement)).thenReturn(1L);
    when(spannerConnection.executeUpdate(largeUpdateStatement)).thenReturn(Integer.MAX_VALUE + 1L);
    try (JdbcStatement statement = new JdbcStatement(connection)) {
        assertThat(statement.executeUpdate(updateStatement)).isEqualTo(1);
        try {
            statement.executeUpdate(largeUpdateStatement);
            fail("missing expected exception");
        } catch (JdbcSqlExceptionImpl e) {
            assertThat(e.getCode()).isEqualTo(Code.OUT_OF_RANGE);
        }
    }
}

Example 15 with Connection

use of com.google.cloud.spanner.connection.Connection in project java-spanner by googleapis.

the class ITAsyncTransactionRetryTest method testAbortWithConcurrentUpdate.

@Test
public void testAbortWithConcurrentUpdate() {
    assumeFalse("concurrent transactions are not supported on the emulator", isUsingEmulator());
    AbortInterceptor interceptor = new AbortInterceptor(0);
    // first insert two test records
    try (ITConnection connection = createConnection()) {
        connection.executeUpdateAsync(Statement.of("INSERT INTO TEST (ID, NAME) VALUES (1, 'test 1')"));
        connection.executeUpdateAsync(Statement.of("INSERT INTO TEST (ID, NAME) VALUES (2, 'test 2')"));
        get(connection.commitAsync());
    }
    // open a new connection and select the two test records
    try (ITConnection connection = createConnection(interceptor, new CountTransactionRetryListener())) {
        // select the test records and consume the entire result set
        try (AsyncResultSet rs = connection.executeQueryAsync(Statement.of("SELECT * FROM TEST ORDER BY ID"))) {
            get(rs.setCallback(executor, resultSet -> {
                while (true) {
                    switch(resultSet.tryNext()) {
                        case DONE:
                            return CallbackResponse.DONE;
                        case NOT_READY:
                            return CallbackResponse.CONTINUE;
                        case OK:
                            break;
                    }
                }
            }));
        }
        // open a new connection and transaction and update one of the test records
        try (ITConnection connection2 = createConnection()) {
            connection2.executeUpdateAsync(Statement.of("UPDATE TEST SET NAME='test updated' WHERE ID=2"));
            get(connection2.commitAsync());
        }
        // now try to do an insert that will abort. The retry should now fail as there has been a
        // concurrent modification
        interceptor.setProbability(1.0);
        interceptor.setOnlyInjectOnce(true);
        try {
            get(connection.executeUpdateAsync(Statement.of("INSERT INTO TEST (ID, NAME) VALUES (3, 'test 3')")));
            fail("Missing expected exception");
        } catch (AbortedDueToConcurrentModificationException e) {
            assertRetryStatistics(1, 1, 0);
        }
    }
}