scwhittle commented on code in PR #40755:
URL: https://github.com/apache/beam/pull/40755#discussion_r4387103267
##########
sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/PartitionMetadataDaoTest.java:
##########
@@ -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<Statement> 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<Statement> 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<Statement> statementCaptor =
ArgumentCaptor.forClass(Statement.class);
+ ArgumentCaptor<ImmutableList<Mutation>> mutations =
+ ArgumentCaptor.forClass(ImmutableList.class);
+ doNothing().when(transaction).buffer(mutations.capture());
+
+
assertNull(pgTransactionContext.updateToScheduled(Collections.singletonList(PARTITION_TOKEN)));
Review Comment:
add test coverage of 0 and multiple tokens as well
##########
sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/PartitionMetadataDao.java:
##########
@@ -681,22 +690,20 @@ private Statement getPartitionsMatchingState(
List<String> 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)")
Review Comment:
This seems to change the behavior for empty composedPartitionTokens array.
Previously it wouldn't check COLUMN_PARTITION_TOKEN but now I think that it
will do ANY() which evaluates to false.
Maybe have a branch on if composedPartitionTokens is empty and have two
statements you bind to?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
