gemini-code-assist[bot] commented on code in PR #36459:
URL: https://github.com/apache/beam/pull/36459#discussion_r2461423385
##########
sdks/java/io/cassandra/src/test/java/org/apache/beam/sdk/io/cassandra/CassandraIOTest.java:
##########
@@ -844,4 +844,378 @@ private static RingRange fromEncodedKey(Metadata
metadata, ByteBuffer... bb) {
/** Simple Cassandra entity used in write tests. */
@Table(name = CASSANDRA_TABLE_WRITE, keyspace = CASSANDRA_KEYSPACE)
static class ScientistWrite extends Scientist {}
+
+ /** Test the quoteIdentifier utility method with various inputs. */
+ @Test
+ public void testQuoteIdentifier() {
+ // Test normal identifiers
+ assertEquals("\"normal_column\"", ReadFn.quoteIdentifier("normal_column"));
+ assertEquals("\"myTable\"", ReadFn.quoteIdentifier("myTable"));
+ assertEquals("\"column123\"", ReadFn.quoteIdentifier("column123"));
+
+ // Test reserved keywords
+ assertEquals("\"true\"", ReadFn.quoteIdentifier("true"));
+ assertEquals("\"key\"", ReadFn.quoteIdentifier("key"));
+ assertEquals("\"select\"", ReadFn.quoteIdentifier("select"));
+ assertEquals("\"from\"", ReadFn.quoteIdentifier("from"));
+ assertEquals("\"where\"", ReadFn.quoteIdentifier("where"));
+ assertEquals("\"table\"", ReadFn.quoteIdentifier("table"));
+ assertEquals("\"keyspace\"", ReadFn.quoteIdentifier("keyspace"));
+
+ // Test identifiers with existing quotes (should be escaped by doubling)
+ assertEquals("\"column\"\"with\"\"quotes\"",
ReadFn.quoteIdentifier("column\"with\"quotes"));
+ assertEquals("\"single\"\"quote\"",
ReadFn.quoteIdentifier("single\"quote"));
+ assertEquals("\"\"\"starts_with_quote\"",
ReadFn.quoteIdentifier("\"starts_with_quote"));
+ assertEquals("\"ends_with_quote\"\"\"",
ReadFn.quoteIdentifier("ends_with_quote\""));
+
+ // Test edge cases
+ assertEquals("\"\"", ReadFn.quoteIdentifier(""));
+ assertNull(ReadFn.quoteIdentifier(null));
+
+ // Test special characters that might be in identifiers
+ assertEquals("\"column with spaces\"", ReadFn.quoteIdentifier("column with
spaces"));
+ assertEquals("\"column-with-dashes\"",
ReadFn.quoteIdentifier("column-with-dashes"));
+ assertEquals("\"column.with.dots\"",
ReadFn.quoteIdentifier("column.with.dots"));
+ }
+
+ /**
+ * Test reading from a table with reserved keyword column names. This
integration test verifies
+ * the complete fix works end-to-end.
+ */
+ @Test
+ public void testReadWithReservedKeywordColumns() throws Exception {
+ String reservedTableName = "reserved_keywords_table";
+
+ // Create table with reserved keyword column names
+ String createTableQuery =
+ String.format(
+ "CREATE TABLE IF NOT EXISTS %s.%s("
+ + "\"true\" text, \"key\" text, \"select\" text, normal_column
text, "
+ + "PRIMARY KEY (\"true\", \"key\")"
+ + ");",
+ CASSANDRA_KEYSPACE, reservedTableName);
+
+ session.execute(createTableQuery);
+
+ // Insert test data with reserved keyword column names
+ String insertQuery1 =
+ String.format(
+ "INSERT INTO %s.%s(\"true\", \"key\", \"select\", normal_column) "
+ + "VALUES ('true_value_1', 'key_value_1', 'select_value_1',
'normal_value_1');",
+ CASSANDRA_KEYSPACE, reservedTableName);
+ session.execute(insertQuery1);
+
+ String insertQuery2 =
+ String.format(
+ "INSERT INTO %s.%s(\"true\", \"key\", \"select\", normal_column) "
+ + "VALUES ('true_value_2', 'key_value_2', 'select_value_2',
'normal_value_2');",
+ CASSANDRA_KEYSPACE, reservedTableName);
+ session.execute(insertQuery2);
+
+ // Flush to ensure data is written
+ flushMemTablesAndRefreshSizeEstimates();
+
+ // Test reading with CassandraIO - this should work with the fix
+ PCollection<ReservedKeywordEntity> output =
+ pipeline.apply(
+ CassandraIO.<ReservedKeywordEntity>read()
+ .withHosts(Collections.singletonList(CASSANDRA_HOST))
+ .withPort(cassandraPort)
+ .withKeyspace(CASSANDRA_KEYSPACE)
+ .withTable(reservedTableName)
+ .withCoder(SerializableCoder.of(ReservedKeywordEntity.class))
+ .withEntity(ReservedKeywordEntity.class));
+
+ // Verify we can read the data successfully
+ PAssert.thatSingleton(output.apply("Count",
Count.globally())).isEqualTo(2L);
+
+ PAssert.that(output)
+ .satisfies(
+ input -> {
+ List<ReservedKeywordEntity> entities = new ArrayList<>();
+ input.forEach(entities::add);
+
+ assertEquals(2, entities.size());
+
+ // Check that data was read correctly
+ boolean foundFirst = false, foundSecond = false;
+ for (ReservedKeywordEntity entity : entities) {
+ if ("true_value_1".equals(entity.trueColumn)) {
+ assertEquals("key_value_1", entity.keyColumn);
+ assertEquals("select_value_1", entity.selectColumn);
+ assertEquals("normal_value_1", entity.normalColumn);
+ foundFirst = true;
+ } else if ("true_value_2".equals(entity.trueColumn)) {
+ assertEquals("key_value_2", entity.keyColumn);
+ assertEquals("select_value_2", entity.selectColumn);
+ assertEquals("normal_value_2", entity.normalColumn);
+ foundSecond = true;
+ }
+ }
+
+ assertTrue("Should find first test record", foundFirst);
+ assertTrue("Should find second test record", foundSecond);
+ return null;
+ });
Review Comment:
The assertions for this test can be simplified and made more readable by
using `PAssert.that(output).containsInAnyOrder(...)`. This avoids manual
iteration and flag-checking, making the test's intent clearer and easier to
maintain. Since `ReservedKeywordEntity` has `equals` and `hashCode`
implemented, this assertion will work as expected.
```java
ReservedKeywordEntity entity1 = new ReservedKeywordEntity();
entity1.trueColumn = "true_value_1";
entity1.keyColumn = "key_value_1";
entity1.selectColumn = "select_value_1";
entity1.normalColumn = "normal_value_1";
ReservedKeywordEntity entity2 = new ReservedKeywordEntity();
entity2.trueColumn = "true_value_2";
entity2.keyColumn = "key_value_2";
entity2.selectColumn = "select_value_2";
entity2.normalColumn = "normal_value_2";
PAssert.that(output).containsInAnyOrder(entity1, entity2);
```
--
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]
