diff --git a/data-prepper-plugins/rds-source/src/main/java/org/opensearch/dataprepper/plugins/source/rds/schema/PostgresSchemaManager.java b/data-prepper-plugins/rds-source/src/main/java/org/opensearch/dataprepper/plugins/source/rds/schema/PostgresSchemaManager.java index 1ea50433ed..ad5decca31 100644 --- a/data-prepper-plugins/rds-source/src/main/java/org/opensearch/dataprepper/plugins/source/rds/schema/PostgresSchemaManager.java +++ b/data-prepper-plugins/rds-source/src/main/java/org/opensearch/dataprepper/plugins/source/rds/schema/PostgresSchemaManager.java @@ -52,7 +52,7 @@ public void createLogicalReplicationSlot(final List tableNames, final St .append(publicationName) .append(" FOR TABLE "); for (int i = 0; i < tableNames.size(); i++) { - createPublicationStatementBuilder.append(tableNames.get(i)); + createPublicationStatementBuilder.append(quoteFullTableName(tableNames.get(i))); if (i < tableNames.size() - 1) { createPublicationStatementBuilder.append(", "); } @@ -293,6 +293,39 @@ Set getEnumColumnsForTable(final Connection connection, final String ful database, schema, table)); } + /** + * Converts a dot-separated fully qualified table name into a form that is safe to embed in generated DDL by + * delimiting each part of the name. PostgreSQL folds undelimited identifiers to lower case and rejects those + * containing characters such as hyphens, so names that are mixed case, hyphenated, or reserved words can only be + * referenced when quoted. + * + * @param fullTableName the fully qualified table name, in {@code database.schema.table} form + * @return the same name with every part delimited, e.g. {@code "My-Db-1"."dbo"."MyTable"} + */ + // Visible for testing + static String quoteFullTableName(final String fullTableName) { + final String[] splits = fullTableName.split("\\."); + final StringBuilder quotedName = new StringBuilder(); + for (int i = 0; i < splits.length; i++) { + if (i > 0) { + quotedName.append("."); + } + quotedName.append(quoteIdentifier(splits[i])); + } + return quotedName.toString(); + } + + /** + * Delimits a single PostgreSQL identifier, escaping any double quote it contains by doubling it. + * + * @param identifier the identifier to delimit + * @return the delimited identifier + */ + // Visible for testing + static String quoteIdentifier(final String identifier) { + return "\"" + identifier.replace("\"", "\"\"") + "\""; + } + private void applyBackoff() { try { Thread.sleep(BACKOFF_IN_MILLIS); diff --git a/data-prepper-plugins/rds-source/src/test/java/org/opensearch/dataprepper/plugins/source/rds/schema/PostgresSchemaManagerTest.java b/data-prepper-plugins/rds-source/src/test/java/org/opensearch/dataprepper/plugins/source/rds/schema/PostgresSchemaManagerTest.java index 91b5ea2f44..859a90dce5 100644 --- a/data-prepper-plugins/rds-source/src/test/java/org/opensearch/dataprepper/plugins/source/rds/schema/PostgresSchemaManagerTest.java +++ b/data-prepper-plugins/rds-source/src/test/java/org/opensearch/dataprepper/plugins/source/rds/schema/PostgresSchemaManagerTest.java @@ -108,7 +108,7 @@ void test_createLogicalReplicationSlot_creates_slot_if_not_exists() throws SQLEx schemaManager.createLogicalReplicationSlot(tableNames, publicationName, slotName); List statements = statementCaptor.getAllValues(); - assertThat(statements.get(0), is("CREATE PUBLICATION " + publicationName + " FOR TABLE " + String.join(", ", tableNames) + ";")); + assertThat(statements.get(0), is("CREATE PUBLICATION " + publicationName + " FOR TABLE \"table1\", \"table2\";")); assertThat(statements.get(1), is("SELECT EXISTS (SELECT 1 FROM pg_replication_slots WHERE slot_name = ?);")); verify(preparedStatement).executeUpdate(); verify(preparedStatement).executeQuery(); @@ -143,7 +143,7 @@ void test_createLogicalReplicationSlot_skip_creation_if_slot_exists() throws SQL schemaManager.createLogicalReplicationSlot(tableNames, publicationName, slotName); List statements = statementCaptor.getAllValues(); - assertThat(statements.get(0), is("CREATE PUBLICATION " + publicationName + " FOR TABLE " + String.join(", ", tableNames) + ";")); + assertThat(statements.get(0), is("CREATE PUBLICATION " + publicationName + " FOR TABLE \"table1\", \"table2\";")); assertThat(statements.get(1), is("SELECT EXISTS (SELECT 1 FROM pg_replication_slots WHERE slot_name = ?);")); verify(preparedStatement).executeUpdate(); verify(preparedStatement).executeQuery(); @@ -151,6 +151,52 @@ void test_createLogicalReplicationSlot_skip_creation_if_slot_exists() throws SQL verify(replicationConnection, never()).createReplicationSlot(); } + @Test + void test_createLogicalReplicationSlot_delimits_identifiers_that_require_quoting() throws SQLException { + final List tableNames = List.of("My-Db-1.dbo.MyTable", "My-Db-1.dbo.my_other_table"); + final String publicationName = "publication1"; + final String slotName = "slot1"; + final PreparedStatement preparedStatement = mock(PreparedStatement.class); + final PGConnection pgConnection = mock(PGConnection.class); + final PGReplicationConnection replicationConnection = mock(PGReplicationConnection.class); + final ResultSet resultSet = mock(ResultSet.class); + + ArgumentCaptor statementCaptor = ArgumentCaptor.forClass(String.class); + + when(connectionManager.getConnection()).thenReturn(connection); + when(connection.prepareStatement(statementCaptor.capture())).thenReturn(preparedStatement); + when(connection.unwrap(PGConnection.class)).thenReturn(pgConnection); + when(preparedStatement.executeQuery()).thenReturn(resultSet); + when(resultSet.next()).thenReturn(true); // Replication slot exists + when(resultSet.getBoolean(1)).thenReturn(true); + when(pgConnection.getReplicationAPI()).thenReturn(replicationConnection); + + schemaManager.createLogicalReplicationSlot(tableNames, publicationName, slotName); + + List statements = statementCaptor.getAllValues(); + assertThat(statements.get(0), is("CREATE PUBLICATION " + publicationName + " FOR TABLE " + + "\"My-Db-1\".\"dbo\".\"MyTable\", \"My-Db-1\".\"dbo\".\"my_other_table\";")); + } + + @Test + void test_quoteFullTableName_delimits_every_part_and_preserves_case() { + assertThat(PostgresSchemaManager.quoteFullTableName("mydb.dbo.MyTable"), + is("\"mydb\".\"dbo\".\"MyTable\"")); + assertThat(PostgresSchemaManager.quoteFullTableName("My-Db-1.dbo.MyTable"), + is("\"My-Db-1\".\"dbo\".\"MyTable\"")); + } + + @Test + void test_quoteFullTableName_when_name_needs_no_quoting_then_only_adds_delimiters() { + assertThat(PostgresSchemaManager.quoteFullTableName("mydb.public.orders"), + is("\"mydb\".\"public\".\"orders\"")); + } + + @Test + void test_quoteIdentifier_escapes_embedded_double_quote() { + assertThat(PostgresSchemaManager.quoteIdentifier("we\"ird"), is("\"we\"\"ird\"")); + } + @Test void test_deleteLogicalReplicationSlot_success() throws SQLException { final String publicationName = UUID.randomUUID().toString();