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

Example 6 with Connection

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

the class JdbcStatementTest method testQueryTimeout.

@Test
public void testQueryTimeout() throws SQLException {
    final String select = "SELECT 1";
    JdbcConnection connection = mock(JdbcConnection.class);
    when(connection.getDialect()).thenReturn(dialect);
    Connection spanner = mock(Connection.class);
    when(connection.getSpannerConnection()).thenReturn(spanner);
    StatementResult result = mock(StatementResult.class);
    when(result.getResultType()).thenReturn(ResultType.RESULT_SET);
    when(result.getResultSet()).thenReturn(mock(com.google.cloud.spanner.ResultSet.class));
    when(spanner.execute(com.google.cloud.spanner.Statement.of(select))).thenReturn(result);
    try (Statement statement = new JdbcStatement(connection)) {
        assertThat(statement.getQueryTimeout()).isEqualTo(0);
        statement.setQueryTimeout(1);
        assertThat(statement.getQueryTimeout()).isEqualTo(1);
        statement.setQueryTimeout(99);
        assertThat(statement.getQueryTimeout()).isEqualTo(99);
        statement.setQueryTimeout(0);
        assertThat(statement.getQueryTimeout()).isEqualTo(0);
    }
    when(spanner.getStatementTimeout(TimeUnit.SECONDS)).thenReturn(1L);
    when(spanner.getStatementTimeout(TimeUnit.MILLISECONDS)).thenReturn(1000L);
    when(spanner.getStatementTimeout(TimeUnit.MICROSECONDS)).thenReturn(1000000L);
    when(spanner.getStatementTimeout(TimeUnit.NANOSECONDS)).thenReturn(1000000000L);
    when(spanner.hasStatementTimeout()).thenReturn(true);
    try (Statement statement = new JdbcStatement(connection)) {
        assertThat(statement.getQueryTimeout()).isEqualTo(0);
        statement.execute(select);
        // statement has no timeout, so it should also not be set on the connection
        verify(spanner, never()).setStatementTimeout(1L, TimeUnit.SECONDS);
    }
    try (Statement statement = new JdbcStatement(connection)) {
        // now set a query timeout that should temporarily applied to the connection
        statement.setQueryTimeout(2);
        statement.execute(select);
        // assert that it is temporarily set to 2 seconds, and then back to the original 1 second
        // value
        verify(spanner).setStatementTimeout(2L, TimeUnit.SECONDS);
        verify(spanner).setStatementTimeout(1L, TimeUnit.SECONDS);
    }
}

Example 7 with Connection

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

the class MutationWriterTest method testWriteMutations_NonAtomic_SucceedsForLargeBatch.

@Test
public void testWriteMutations_NonAtomic_SucceedsForLargeBatch() throws Exception {
    // 6 == 2 mutations per batch, as we have 2 columns + 1 indexed column.
    System.setProperty("copy_in_mutation_limit", "6");
    try {
        Map<String, TypeCode> tableColumns = ImmutableMap.of("number", TypeCode.INT64, "name", TypeCode.STRING);
        CSVFormat format = CSVFormat.POSTGRESQL_TEXT.builder().setHeader(tableColumns.keySet().toArray(new String[0])).build();
        Connection connection = mock(Connection.class);
        DatabaseClient databaseClient = mock(DatabaseClient.class);
        MutationWriter mutationWriter = new MutationWriter(CopyTransactionMode.ImplicitNonAtomic, connection, databaseClient, "numbers", tableColumns, /* indexedColumnsCount = */
        1, format, false);
        mutationWriter.addCopyData("1\t\"One\"\n2\t\"Two\"\n3\t\"Three\"\n4\t\"Four\"\n5\t\"Five\"\n".getBytes(StandardCharsets.UTF_8));
        mutationWriter.close();
        long updateCount = mutationWriter.call();
        assertEquals(5L, updateCount);
        verify(databaseClient, times(3)).write(anyIterable());
    } finally {
        System.getProperties().remove("copy_in_mutation_limit");
    }
}

Example 8 with Connection

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

the class MutationWriterTest method testWritePartials.

@Test
public void testWritePartials() throws Exception {
    Map<String, TypeCode> tableColumns = ImmutableMap.of("number", TypeCode.INT64, "name", TypeCode.STRING);
    CSVFormat format = CSVFormat.POSTGRESQL_TEXT.builder().setHeader(tableColumns.keySet().toArray(new String[0])).build();
    Connection connection = mock(Connection.class);
    DatabaseClient databaseClient = mock(DatabaseClient.class);
    MutationWriter mutationWriter = new MutationWriter(CopyTransactionMode.ImplicitNonAtomic, connection, databaseClient, "numbers", tableColumns, /* indexedColumnsCount = */
    1, format, false);
    mutationWriter.addCopyData("1\t\"One\"\n2\t\"Two\"\n3\t\"Three\"\n4\t\"Four\"\n5\t\"Five\"\n".getBytes(StandardCharsets.UTF_8));
    mutationWriter.close();
    ExecutorService executor = Executors.newFixedThreadPool(2);
    executor.submit(() -> {
        mutationWriter.addCopyData("1\t\"One\"\n".getBytes(StandardCharsets.UTF_8));
        mutationWriter.addCopyData("2\t\"Two".getBytes(StandardCharsets.UTF_8));
        mutationWriter.addCopyData("\"".getBytes(StandardCharsets.UTF_8));
        mutationWriter.addCopyData("\n3\t".getBytes(StandardCharsets.UTF_8));
        mutationWriter.addCopyData("\"Three\"\n4\t\"Four\"\n5\t".getBytes(StandardCharsets.UTF_8));
        mutationWriter.addCopyData("\"Five\"\n".getBytes(StandardCharsets.UTF_8));
        mutationWriter.close();
        return null;
    });
    Future<Long> updateCount = executor.submit(mutationWriter);
    assertEquals(5L, updateCount.get().longValue());
    List<Mutation> expectedMutations = ImmutableList.of(Mutation.newInsertBuilder("numbers").set("number").to(1L).set("name").to("One").build(), Mutation.newInsertBuilder("numbers").set("number").to(2L).set("name").to("Two").build(), Mutation.newInsertBuilder("numbers").set("number").to(3L).set("name").to("Three").build(), Mutation.newInsertBuilder("numbers").set("number").to(4L).set("name").to("Four").build(), Mutation.newInsertBuilder("numbers").set("number").to(5L).set("name").to("Five").build());
    verify(databaseClient).write(expectedMutations);
    executor.shutdown();
}

Example 9 with Connection

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

the class MutationWriterTest method testWriteMutations.

@Test
public void testWriteMutations() throws Exception {
    Map<String, TypeCode> tableColumns = ImmutableMap.of("number", TypeCode.INT64, "name", TypeCode.STRING);
    CSVFormat format = CSVFormat.POSTGRESQL_TEXT.builder().setHeader(tableColumns.keySet().toArray(new String[0])).build();
    Connection connection = mock(Connection.class);
    DatabaseClient databaseClient = mock(DatabaseClient.class);
    MutationWriter mutationWriter = new MutationWriter(CopyTransactionMode.ImplicitAtomic, connection, databaseClient, "numbers", tableColumns, /* indexedColumnsCount = */
    1, format, false);
    mutationWriter.addCopyData("1\t\"One\"\n2\t\"Two\"\n".getBytes(StandardCharsets.UTF_8));
    mutationWriter.close();
    long updateCount = mutationWriter.call();
    assertEquals(2L, updateCount);
    List<Mutation> expectedMutations = ImmutableList.of(Mutation.newInsertBuilder("numbers").set("number").to(1L).set("name").to("One").build(), Mutation.newInsertBuilder("numbers").set("number").to(2L).set("name").to("Two").build());
    verify(databaseClient).write(expectedMutations);
}

Example 10 with Connection

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

the class MutationWriterTest method testWriteMutations_FailsForLargeBatch.

@Test
public void testWriteMutations_FailsForLargeBatch() throws Exception {
    System.setProperty("copy_in_mutation_limit", "1");
    try {
        Map<String, TypeCode> tableColumns = ImmutableMap.of("number", TypeCode.INT64, "name", TypeCode.STRING);
        CSVFormat format = CSVFormat.POSTGRESQL_TEXT.builder().setHeader(tableColumns.keySet().toArray(new String[0])).build();
        Connection connection = mock(Connection.class);
        MutationWriter mutationWriter = new MutationWriter(CopyTransactionMode.ImplicitAtomic, connection, "numbers", tableColumns, /* indexedColumnsCount = */
        1, format, false);
        mutationWriter.addCopyData("1\t\"One\"\n2\t\"Two\"\n".getBytes(StandardCharsets.UTF_8));
        mutationWriter.close();
        SpannerException exception = assertThrows(SpannerException.class, mutationWriter::call);
        assertTrue(exception.getMessage().contains("Record count: 2 has exceeded the limit: 1"));
        verify(connection, never()).write(anyIterable());
    } finally {
        System.getProperties().remove("copy_in_mutation_limit");
    }
}