diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/PartitionMetadataDao.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/PartitionMetadataDao.java index a60bd5c2143f..7b0825cf615a 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/PartitionMetadataDao.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/PartitionMetadataDao.java @@ -50,7 +50,6 @@ import java.util.List; import java.util.Map; import java.util.function.Function; -import java.util.stream.Collectors; import javax.annotation.Nullable; import org.apache.beam.sdk.io.gcp.spanner.changestreams.model.PartitionMetadata; import org.apache.beam.sdk.io.gcp.spanner.changestreams.model.PartitionMetadata.State; @@ -83,13 +82,23 @@ public class PartitionMetadataDao { * @return true if the table exists, false if the table does not exist. */ public boolean tableExists() { - final String checkTableExistsStmt = - "SELECT t.table_name FROM information_schema.tables AS t " - + "WHERE t.table_catalog = '' AND " - + "t.table_schema = '' AND " - + "t.table_name = '" - + metadataTableName - + "'"; + final String checkTableExistsStmt; + if (this.isPostgres()) { + checkTableExistsStmt = + "SELECT t.table_name FROM information_schema.tables AS t " + + "WHERE t.table_schema = 'public' AND " + + "t.table_name = '" + + metadataTableName + + "'"; + } else { + checkTableExistsStmt = + "SELECT t.table_name FROM information_schema.tables AS t " + + "WHERE t.table_catalog = '' AND " + + "t.table_schema = '' AND " + + "t.table_name = '" + + metadataTableName + + "'"; + } try (ResultSet queryResultSet = databaseClient .singleUseReadOnlyTransaction() @@ -681,22 +690,20 @@ private Statement getPartitionsMatchingState( List composedPartitionTokens, State state) { Statement statement; if (this.dialect == Dialect.POSTGRESQL) { - StringBuilder sqlStringBuilder = - new StringBuilder("SELECT * FROM \"" + metadataTableName + "\""); - sqlStringBuilder.append(" WHERE \""); - sqlStringBuilder.append(COLUMN_STATE + "\" = " + "'" + state.toString() + "'"); - if (!composedPartitionTokens.isEmpty()) { - sqlStringBuilder.append(" AND \""); - sqlStringBuilder.append(COLUMN_PARTITION_TOKEN); - sqlStringBuilder.append("\""); - sqlStringBuilder.append(" = ANY (Array["); - sqlStringBuilder.append( - composedPartitionTokens.stream() - .map(s -> "'" + s + "'") - .collect(Collectors.joining(","))); - sqlStringBuilder.append("])"); - } - statement = Statement.newBuilder(sqlStringBuilder.toString()).build(); + statement = + Statement.newBuilder( + "SELECT * FROM \"" + + metadataTableName + + "\" WHERE \"" + + COLUMN_STATE + + "\" = $1 AND \"" + + COLUMN_PARTITION_TOKEN + + "\" = ANY ($2)") + .bind("p1") + .to(state.toString()) + .bind("p2") + .to(Value.stringArray(new ArrayList<>(composedPartitionTokens))) + .build(); } else { statement = Statement.newBuilder( diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/PartitionMetadataDaoTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/PartitionMetadataDaoTest.java index 79fe01a2b7d8..c02beca2d4a6 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/PartitionMetadataDaoTest.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/PartitionMetadataDaoTest.java @@ -18,8 +18,10 @@ package org.apache.beam.sdk.io.gcp.spanner.changestreams.dao; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.doNothing; import static org.mockito.Mockito.mock; @@ -31,7 +33,9 @@ import com.google.cloud.spanner.DatabaseClient; import com.google.cloud.spanner.Dialect; import com.google.cloud.spanner.Mutation; +import com.google.cloud.spanner.ReadOnlyTransaction; import com.google.cloud.spanner.ResultSet; +import com.google.cloud.spanner.Statement; import com.google.cloud.spanner.Struct; import com.google.cloud.spanner.TransactionContext; import com.google.cloud.spanner.TransactionRunner; @@ -406,4 +410,81 @@ public void testExtractTvfName() { assertEquals("tvf", PartitionMetadataDao.extractTvfName("partitionToken123#tvf")); assertEquals("", PartitionMetadataDao.extractTvfName("partitionToken123")); } + + @Test + public void testTableExistsGoogleSql() { + ReadOnlyTransaction readOnlyTx = mock(ReadOnlyTransaction.class); + ResultSet resultSet = mock(ResultSet.class); + when(databaseClient.singleUseReadOnlyTransaction()).thenReturn(readOnlyTx); + when(readOnlyTx.executeQuery(any(Statement.class), any())).thenReturn(resultSet); + when(resultSet.next()).thenReturn(true); + + assertTrue(partitionMetadataDao.tableExists()); + + ArgumentCaptor statementCaptor = ArgumentCaptor.forClass(Statement.class); + verify(readOnlyTx).executeQuery(statementCaptor.capture(), any()); + assertEquals( + "SELECT t.table_name FROM information_schema.tables AS t " + + "WHERE t.table_catalog = '' AND " + + "t.table_schema = '' AND " + + "t.table_name = '" + + METADATA_TABLE_NAME + + "'", + statementCaptor.getValue().getSql()); + } + + @Test + public void testTableExistsPostgres() { + ReadOnlyTransaction readOnlyTx = mock(ReadOnlyTransaction.class); + ResultSet resultSet = mock(ResultSet.class); + when(databaseClient.singleUseReadOnlyTransaction()).thenReturn(readOnlyTx); + when(readOnlyTx.executeQuery(any(Statement.class), any())).thenReturn(resultSet); + when(resultSet.next()).thenReturn(true).thenReturn(false); + + PartitionMetadataDao pgDao = + new PartitionMetadataDao(METADATA_TABLE_NAME, databaseClient, Dialect.POSTGRESQL); + assertTrue(pgDao.tableExists()); + assertFalse(pgDao.tableExists()); + + ArgumentCaptor statementCaptor = ArgumentCaptor.forClass(Statement.class); + verify(readOnlyTx, times(2)).executeQuery(statementCaptor.capture(), any()); + assertEquals( + "SELECT t.table_name FROM information_schema.tables AS t " + + "WHERE t.table_schema = 'public' AND " + + "t.table_name = '" + + METADATA_TABLE_NAME + + "'", + statementCaptor.getValue().getSql()); + } + + @Test + public void testInTransactionContextUpdateToScheduledPostgres() { + PartitionMetadataDao.InTransactionContext pgTransactionContext = + new PartitionMetadataDao.InTransactionContext( + METADATA_TABLE_NAME, transaction, Dialect.POSTGRESQL); + ResultSet resultSet = mock(ResultSet.class); + when(transaction.executeQuery(any(), any())).thenReturn(resultSet); + when(resultSet.next()).thenReturn(true).thenReturn(false); + when(resultSet.getString(any())).thenReturn(PARTITION_TOKEN); + + ArgumentCaptor statementCaptor = ArgumentCaptor.forClass(Statement.class); + ArgumentCaptor> mutations = + ArgumentCaptor.forClass(ImmutableList.class); + doNothing().when(transaction).buffer(mutations.capture()); + + assertNull(pgTransactionContext.updateToScheduled(Collections.singletonList(PARTITION_TOKEN))); + + verify(transaction).executeQuery(statementCaptor.capture(), any()); + Statement captured = statementCaptor.getValue(); + assertEquals( + "SELECT * FROM \"" + + METADATA_TABLE_NAME + + "\" WHERE \"State\" = $1 AND \"PartitionToken\" = ANY ($2)", + captured.getSql()); + assertEquals(State.CREATED.toString(), captured.getParameters().get("p1").getString()); + assertEquals( + Collections.singletonList(PARTITION_TOKEN), + captured.getParameters().get("p2").getStringArray()); + assertEquals(1, mutations.getValue().size()); + } }