Skip to content

Commit f63eea0

Browse files
committed
fix(postgres): schema-qualify the primary-key subquery in getTableColumns
The PK subquery joined tc.constraint_name = ku.constraint_name and filtered on ku.table_name alone, with no schema predicate anywhere. Postgres auto-names primary keys "<table>_pkey", so two schemas holding a same-named table collide by construction — the join cross-matches them. Measured against two schemas each with an `orders` table (s_a's PK is `other`, s_b's is `name`), asking for s_b.orders: column_name | is_primary_key name | t name | t <- duplicated row other | t <- false positive: `other` is s_a's PK, not s_b's other | t <- duplicated row Ground truth from pg_index: s_b.orders has exactly one PK column, `name`. So the provider reported every column as a primary key and returned each one twice — duplicate ColumnInfo entries, and a table with no real single PK. Adding tc.table_schema = ku.table_schema and ku.table_schema = ? returns the correct `name | t`, `other | f`. This predates the search_path change and was mostly latent while the provider only ever read `public`: a collision needed some other schema to hold a same-named table. Now that it targets whatever the session resolves to, that is no longer the exception — staging/marts/public each holding `orders` or `customers` is the normal shape of the dbt warehouse this change was made for. So the earlier commit does not introduce this bug, but it does move it from unlikely to expected, which is why it is fixed on the same branch. Also hoists resolveSchema() out of the argument list: this method runs once per table, and each call costs a round-trip. Tests: 12/12. A mocked ResultSet cannot exercise SQL semantics, so the new test asserts the two schema predicates are present and that every placeholder is bound — enough to stop the qualification being dropped again. The correctness evidence is the measurement above, run against Postgres 18.
1 parent 2a23bc4 commit f63eea0

2 files changed

Lines changed: 64 additions & 4 deletions

File tree

backend/src/main/java/com/dbaagent/provider/postgres/PostgresIntrospectionProvider.java

Lines changed: 32 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -156,24 +156,52 @@ SELECT p.proname as name, pg_get_functiondef(p.oid) as definition
156156
public List<ColumnInfo> getTableColumns(Connection connection, String database, String tableName) throws SQLException {
157157
List<ColumnInfo> columns = new ArrayList<>();
158158

159+
// The PK subquery must be schema-qualified on BOTH the constraint join and
160+
// the table lookup. Without it, two schemas holding a same-named table
161+
// collide by construction: Postgres auto-names primary keys
162+
// "<table>_pkey", so `tc.constraint_name = ku.constraint_name` alone
163+
// cross-joins the schemas.
164+
//
165+
// Measured against two schemas each holding an `orders` table, asking for
166+
// s_b.orders whose only PK is `name`:
167+
// name | t
168+
// name | t <- duplicated row
169+
// other | t <- false positive; `other` is s_a's PK, not s_b's
170+
// other | t <- duplicated row
171+
// i.e. every column reported as a primary key, and each one twice. With the
172+
// predicates below the same query returns `name | t`, `other | f`.
173+
//
174+
// The bug predates the search_path change but was mostly latent while this
175+
// provider only ever read `public`. Now that it targets whatever the session
176+
// resolves to, same-named tables across schemas — the normal shape of a dbt
177+
// warehouse (staging/marts/public each with `orders`) — are the expected
178+
// case rather than the exception.
159179
String query = """
160180
SELECT c.column_name, c.data_type, c.is_nullable, c.column_default,
161181
CASE WHEN pk.column_name IS NOT NULL THEN true ELSE false END as is_primary_key
162182
FROM information_schema.columns c
163183
LEFT JOIN (
164184
SELECT ku.column_name
165185
FROM information_schema.table_constraints tc
166-
JOIN information_schema.key_column_usage ku ON tc.constraint_name = ku.constraint_name
167-
WHERE tc.constraint_type = 'PRIMARY KEY' AND ku.table_name = ?
186+
JOIN information_schema.key_column_usage ku
187+
ON tc.constraint_name = ku.constraint_name
188+
AND tc.table_schema = ku.table_schema
189+
WHERE tc.constraint_type = 'PRIMARY KEY'
190+
AND ku.table_name = ? AND ku.table_schema = ?
168191
) pk ON c.column_name = pk.column_name
169192
WHERE c.table_name = ? AND c.table_schema = ?
170193
ORDER BY c.ordinal_position
171194
""";
172195

196+
// Resolved once: this method runs per table, and resolveSchema() costs a
197+
// round-trip each call.
198+
String schema = resolveSchema(connection);
199+
173200
try (PreparedStatement stmt = connection.prepareStatement(query)) {
174201
stmt.setString(1, tableName);
175-
stmt.setString(2, tableName);
176-
stmt.setString(3, resolveSchema(connection));
202+
stmt.setString(2, schema);
203+
stmt.setString(3, tableName);
204+
stmt.setString(4, schema);
177205
try (ResultSet rs = stmt.executeQuery()) {
178206
while (rs.next()) {
179207
ColumnInfo col = new ColumnInfo();

backend/src/test/java/com/dbaagent/provider/postgres/PostgresIntrospectionProviderTest.java

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -198,6 +198,38 @@ void getTableStats_returnsStats() throws SQLException {
198198
assertEquals(122880L, stats.getSizeBytes());
199199
}
200200

201+
@Test
202+
void getTableColumns_qualifiesThePrimaryKeySubqueryBySchema() throws SQLException {
203+
// Postgres auto-names primary keys "<table>_pkey", so joining
204+
// tc.constraint_name = ku.constraint_name WITHOUT a schema predicate
205+
// cross-joins any two schemas holding a same-named table. Measured against
206+
// two `orders` tables (s_a PK `other`, s_b PK `name`), asking for s_b:
207+
// name|t name|t other|t other|t
208+
// — every column a primary key, and each row duplicated. Schema-qualified
209+
// it returns name|t, other|f.
210+
//
211+
// A mocked ResultSet cannot exercise SQL semantics, so this asserts the
212+
// predicates are present and every placeholder is bound — enough to stop
213+
// the qualification being dropped again.
214+
ArgumentCaptor<String> sqlCaptor = ArgumentCaptor.forClass(String.class);
215+
when(connection.prepareStatement(anyString())).thenReturn(preparedStatement);
216+
when(preparedStatement.executeQuery()).thenReturn(resultSet);
217+
when(resultSet.next()).thenReturn(false);
218+
219+
provider.getTableColumns(connection, "public", "orders");
220+
221+
verify(connection).prepareStatement(sqlCaptor.capture());
222+
String sql = sqlCaptor.getValue();
223+
224+
assertTrue(sql.contains("tc.table_schema = ku.table_schema"),
225+
"the constraint join must be schema-qualified, or <table>_pkey collides across schemas");
226+
assertTrue(sql.contains("ku.table_schema = ?"),
227+
"the PK lookup must be restricted to the target schema");
228+
229+
int placeholders = (int) sql.chars().filter(c -> c == '?').count();
230+
verify(preparedStatement, times(placeholders)).setString(anyInt(), anyString());
231+
}
232+
201233
@Test
202234
void getTableStats_bindsEveryPlaceholderInTheStatsQuery() throws SQLException {
203235
// The stats query carried NINE `?` placeholders (the two size subtractions use

0 commit comments

Comments
 (0)