Skip to content

Commit b900749

Browse files
suddendustsuddendust
committed
Fix failing test case
1 parent 819fa2b commit b900749

1 file changed

Lines changed: 66 additions & 60 deletions

File tree

‎document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresMetadataFetcherTest.java‎

Lines changed: 66 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -29,52 +29,56 @@ class PostgresMetadataFetcherTest {
2929

3030
@Mock private PostgresClient client;
3131
@Mock private Connection connection;
32-
@Mock private PreparedStatement columnsPreparedStatement;
33-
@Mock private PreparedStatement pkPreparedStatement;
34-
@Mock private ResultSet columnsResultSet;
35-
@Mock private ResultSet pkResultSet;
32+
@Mock private PreparedStatement preparedStatement;
33+
@Mock private ResultSet resultSet;
3634

3735
private PostgresMetadataFetcher fetcher;
3836

3937
private static final String DISCOVERY_SQL =
40-
"SELECT column_name, udt_name, is_nullable "
41-
+ "FROM information_schema.columns "
42-
+ "WHERE table_schema = 'public' AND table_name = ?";
43-
44-
private static final String PRIMARY_KEY_SQL =
45-
"SELECT a.attname as column_name "
46-
+ "FROM pg_index i "
47-
+ "JOIN pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = ANY(i.indkey) "
48-
+ "WHERE i.indrelid = ?::regclass AND i.indisprimary";
38+
"SELECT a.attname AS column_name, "
39+
+ "t.typname AS udt_name, "
40+
+ "NOT a.attnotnull AS is_nullable, "
41+
+ "COALESCE(pk.is_pk, false) AS is_primary_key "
42+
+ "FROM pg_attribute a "
43+
+ "JOIN pg_class c ON a.attrelid = c.oid "
44+
+ "JOIN pg_namespace n ON c.relnamespace = n.oid "
45+
+ "JOIN pg_type t ON a.atttypid = t.oid "
46+
+ "LEFT JOIN ( "
47+
+ " SELECT a2.attname, true AS is_pk "
48+
+ " FROM pg_index i "
49+
+ " JOIN pg_attribute a2 ON a2.attrelid = i.indrelid AND a2.attnum = ANY(i.indkey) "
50+
+ " WHERE i.indisprimary "
51+
+ ") pk ON pk.attname = a.attname "
52+
+ "WHERE c.relname = ? "
53+
+ "AND n.nspname = 'public' "
54+
+ "AND a.attnum > 0 "
55+
+ "AND NOT a.attisdropped";
4956

5057
@BeforeEach
5158
void setUp() throws SQLException {
5259
when(client.getPooledConnection()).thenReturn(connection);
53-
when(connection.prepareStatement(DISCOVERY_SQL)).thenReturn(columnsPreparedStatement);
54-
when(connection.prepareStatement(PRIMARY_KEY_SQL)).thenReturn(pkPreparedStatement);
55-
when(columnsPreparedStatement.executeQuery()).thenReturn(columnsResultSet);
56-
when(pkPreparedStatement.executeQuery()).thenReturn(pkResultSet);
57-
// Default: no primary keys
58-
when(pkResultSet.next()).thenReturn(false);
60+
when(connection.prepareStatement(DISCOVERY_SQL)).thenReturn(preparedStatement);
61+
when(preparedStatement.executeQuery()).thenReturn(resultSet);
5962
fetcher = new PostgresMetadataFetcher(client);
6063
}
6164

6265
@Test
6366
void fetchReturnsEmptyMapForTableWithNoColumns() throws SQLException {
64-
when(columnsResultSet.next()).thenReturn(false);
67+
when(resultSet.next()).thenReturn(false);
6568

6669
Map<String, PostgresColumnMetadata> result = fetcher.fetch(TEST_TABLE);
6770

6871
assertTrue(result.isEmpty());
69-
verify(columnsPreparedStatement).setString(1, TEST_TABLE);
72+
verify(preparedStatement).setString(1, TEST_TABLE);
7073
}
7174

7275
@Test
7376
void fetchReturnsSingleColumn() throws SQLException {
74-
when(columnsResultSet.next()).thenReturn(true, false);
75-
when(columnsResultSet.getString("column_name")).thenReturn("id");
76-
when(columnsResultSet.getString("udt_name")).thenReturn("int4");
77-
when(columnsResultSet.getString("is_nullable")).thenReturn("NO");
77+
when(resultSet.next()).thenReturn(true, false);
78+
when(resultSet.getString("column_name")).thenReturn("id");
79+
when(resultSet.getString("udt_name")).thenReturn("int4");
80+
when(resultSet.getBoolean("is_nullable")).thenReturn(false);
81+
when(resultSet.getBoolean("is_primary_key")).thenReturn(false);
7882

7983
Map<String, PostgresColumnMetadata> result = fetcher.fetch(TEST_TABLE);
8084

@@ -89,10 +93,11 @@ void fetchReturnsSingleColumn() throws SQLException {
8993

9094
@Test
9195
void fetchReturnsMultipleColumns() throws SQLException {
92-
when(columnsResultSet.next()).thenReturn(true, true, true, false);
93-
when(columnsResultSet.getString("column_name")).thenReturn("id", "name", "price");
94-
when(columnsResultSet.getString("udt_name")).thenReturn("int8", "text", "float8");
95-
when(columnsResultSet.getString("is_nullable")).thenReturn("NO", "YES", "YES");
96+
when(resultSet.next()).thenReturn(true, true, true, false);
97+
when(resultSet.getString("column_name")).thenReturn("id", "name", "price");
98+
when(resultSet.getString("udt_name")).thenReturn("int8", "text", "float8");
99+
when(resultSet.getBoolean("is_nullable")).thenReturn(false, true, true);
100+
when(resultSet.getBoolean("is_primary_key")).thenReturn(false, false, false);
96101

97102
Map<String, PostgresColumnMetadata> result = fetcher.fetch(TEST_TABLE);
98103

@@ -119,7 +124,7 @@ void fetchReturnsMultipleColumns() throws SQLException {
119124

120125
@Test
121126
void fetchMapsInt4ToInteger() throws SQLException {
122-
setupSingleColumnResult("col", "int4", "NO");
127+
setupSingleColumnResult("col", "int4", false);
123128

124129
PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col");
125130

@@ -129,7 +134,7 @@ void fetchMapsInt4ToInteger() throws SQLException {
129134

130135
@Test
131136
void fetchMapsInt2ToInteger() throws SQLException {
132-
setupSingleColumnResult("col", "int2", "NO");
137+
setupSingleColumnResult("col", "int2", false);
133138

134139
PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col");
135140

@@ -139,7 +144,7 @@ void fetchMapsInt2ToInteger() throws SQLException {
139144

140145
@Test
141146
void fetchMapsInt8ToLong() throws SQLException {
142-
setupSingleColumnResult("col", "int8", "NO");
147+
setupSingleColumnResult("col", "int8", false);
143148

144149
PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col");
145150

@@ -149,7 +154,7 @@ void fetchMapsInt8ToLong() throws SQLException {
149154

150155
@Test
151156
void fetchMapsFloat4ToFloat() throws SQLException {
152-
setupSingleColumnResult("col", "float4", "NO");
157+
setupSingleColumnResult("col", "float4", false);
153158

154159
PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col");
155160

@@ -159,7 +164,7 @@ void fetchMapsFloat4ToFloat() throws SQLException {
159164

160165
@Test
161166
void fetchMapsFloat8ToDouble() throws SQLException {
162-
setupSingleColumnResult("col", "float8", "NO");
167+
setupSingleColumnResult("col", "float8", false);
163168

164169
PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col");
165170

@@ -169,7 +174,7 @@ void fetchMapsFloat8ToDouble() throws SQLException {
169174

170175
@Test
171176
void fetchMapsNumericToDouble() throws SQLException {
172-
setupSingleColumnResult("col", "numeric", "NO");
177+
setupSingleColumnResult("col", "numeric", false);
173178

174179
PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col");
175180

@@ -179,7 +184,7 @@ void fetchMapsNumericToDouble() throws SQLException {
179184

180185
@Test
181186
void fetchMapsBoolToBoolean() throws SQLException {
182-
setupSingleColumnResult("col", "bool", "NO");
187+
setupSingleColumnResult("col", "bool", false);
183188

184189
PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col");
185190

@@ -189,7 +194,7 @@ void fetchMapsBoolToBoolean() throws SQLException {
189194

190195
@Test
191196
void fetchMapsTextToString() throws SQLException {
192-
setupSingleColumnResult("col", "text", "NO");
197+
setupSingleColumnResult("col", "text", false);
193198

194199
PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col");
195200

@@ -199,7 +204,7 @@ void fetchMapsTextToString() throws SQLException {
199204

200205
@Test
201206
void fetchMapsVarcharToString() throws SQLException {
202-
setupSingleColumnResult("col", "varchar", "NO");
207+
setupSingleColumnResult("col", "varchar", false);
203208

204209
PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col");
205210

@@ -209,7 +214,7 @@ void fetchMapsVarcharToString() throws SQLException {
209214

210215
@Test
211216
void fetchMapsBpcharToString() throws SQLException {
212-
setupSingleColumnResult("col", "bpchar", "NO");
217+
setupSingleColumnResult("col", "bpchar", false);
213218

214219
PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col");
215220

@@ -219,7 +224,7 @@ void fetchMapsBpcharToString() throws SQLException {
219224

220225
@Test
221226
void fetchMapsUuidToString() throws SQLException {
222-
setupSingleColumnResult("col", "uuid", "NO");
227+
setupSingleColumnResult("col", "uuid", false);
223228

224229
PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col");
225230

@@ -229,7 +234,7 @@ void fetchMapsUuidToString() throws SQLException {
229234

230235
@Test
231236
void fetchMapsJsonbToJson() throws SQLException {
232-
setupSingleColumnResult("col", "jsonb", "NO");
237+
setupSingleColumnResult("col", "jsonb", false);
233238

234239
PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col");
235240

@@ -239,7 +244,7 @@ void fetchMapsJsonbToJson() throws SQLException {
239244

240245
@Test
241246
void fetchMapsTimestamptzToTimestamptz() throws SQLException {
242-
setupSingleColumnResult("col", "timestamptz", "NO");
247+
setupSingleColumnResult("col", "timestamptz", false);
243248

244249
PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col");
245250

@@ -249,7 +254,7 @@ void fetchMapsTimestamptzToTimestamptz() throws SQLException {
249254

250255
@Test
251256
void fetchMapsDateToDate() throws SQLException {
252-
setupSingleColumnResult("col", "date", "NO");
257+
setupSingleColumnResult("col", "date", false);
253258

254259
PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col");
255260

@@ -259,7 +264,7 @@ void fetchMapsDateToDate() throws SQLException {
259264

260265
@Test
261266
void fetchMapsUnknownTypeToUnspecified() throws SQLException {
262-
setupSingleColumnResult("col", "unknown_type", "NO");
267+
setupSingleColumnResult("col", "unknown_type", false);
263268

264269
PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col");
265270

@@ -269,7 +274,7 @@ void fetchMapsUnknownTypeToUnspecified() throws SQLException {
269274

270275
@Test
271276
void fetchMapsNullUdtNameToUnspecified() throws SQLException {
272-
setupSingleColumnResult("col", null, "NO");
277+
setupSingleColumnResult("col", null, false);
273278

274279
PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col");
275280

@@ -279,7 +284,7 @@ void fetchMapsNullUdtNameToUnspecified() throws SQLException {
279284

280285
@Test
281286
void fetchHandlesNullableColumn() throws SQLException {
282-
setupSingleColumnResult("col", "text", "YES");
287+
setupSingleColumnResult("col", "text", true);
283288

284289
PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col");
285290

@@ -288,7 +293,7 @@ void fetchHandlesNullableColumn() throws SQLException {
288293

289294
@Test
290295
void fetchHandlesNonNullableColumn() throws SQLException {
291-
setupSingleColumnResult("col", "text", "NO");
296+
setupSingleColumnResult("col", "text", false);
292297

293298
PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col");
294299

@@ -297,7 +302,7 @@ void fetchHandlesNonNullableColumn() throws SQLException {
297302

298303
@Test
299304
void fetchHandlesCaseInsensitiveUdtName() throws SQLException {
300-
setupSingleColumnResult("col", "INT4", "NO");
305+
setupSingleColumnResult("col", "INT4", false);
301306

302307
PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col");
303308

@@ -307,7 +312,7 @@ void fetchHandlesCaseInsensitiveUdtName() throws SQLException {
307312

308313
@Test
309314
void fetchThrowsRuntimeExceptionOnSqlException() throws SQLException {
310-
when(columnsPreparedStatement.executeQuery()).thenThrow(new SQLException("Connection failed"));
315+
when(preparedStatement.executeQuery()).thenThrow(new SQLException("Connection failed"));
311316

312317
RuntimeException exception =
313318
assertThrows(RuntimeException.class, () -> fetcher.fetch(TEST_TABLE));
@@ -318,7 +323,7 @@ void fetchThrowsRuntimeExceptionOnSqlException() throws SQLException {
318323

319324
@Test
320325
void fetchMapsTextArrayToStringArray() throws SQLException {
321-
setupSingleColumnResult("col", "_text", "NO");
326+
setupSingleColumnResult("col", "_text", false);
322327

323328
PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col");
324329

@@ -329,7 +334,7 @@ void fetchMapsTextArrayToStringArray() throws SQLException {
329334

330335
@Test
331336
void fetchMapsInt4ArrayToIntegerArray() throws SQLException {
332-
setupSingleColumnResult("col", "_int4", "NO");
337+
setupSingleColumnResult("col", "_int4", false);
333338

334339
PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col");
335340

@@ -340,7 +345,7 @@ void fetchMapsInt4ArrayToIntegerArray() throws SQLException {
340345

341346
@Test
342347
void fetchMapsInt8ArrayToLongArray() throws SQLException {
343-
setupSingleColumnResult("col", "_int8", "NO");
348+
setupSingleColumnResult("col", "_int8", false);
344349

345350
PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col");
346351

@@ -351,7 +356,7 @@ void fetchMapsInt8ArrayToLongArray() throws SQLException {
351356

352357
@Test
353358
void fetchMapsFloat8ArrayToDoubleArray() throws SQLException {
354-
setupSingleColumnResult("col", "_float8", "NO");
359+
setupSingleColumnResult("col", "_float8", false);
355360

356361
PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col");
357362

@@ -362,7 +367,7 @@ void fetchMapsFloat8ArrayToDoubleArray() throws SQLException {
362367

363368
@Test
364369
void fetchMapsBoolArrayToBooleanArray() throws SQLException {
365-
setupSingleColumnResult("col", "_bool", "NO");
370+
setupSingleColumnResult("col", "_bool", false);
366371

367372
PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col");
368373

@@ -373,18 +378,19 @@ void fetchMapsBoolArrayToBooleanArray() throws SQLException {
373378

374379
@Test
375380
void fetchReturnsIsArrayFalseForNonArrayTypes() throws SQLException {
376-
setupSingleColumnResult("col", "text", "NO");
381+
setupSingleColumnResult("col", "text", false);
377382

378383
PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col");
379384

380385
assertFalse(meta.isArray());
381386
}
382387

383-
private void setupSingleColumnResult(String colName, String udtName, String isNullable)
388+
private void setupSingleColumnResult(String colName, String udtName, boolean isNullable)
384389
throws SQLException {
385-
when(columnsResultSet.next()).thenReturn(true, false);
386-
when(columnsResultSet.getString("column_name")).thenReturn(colName);
387-
when(columnsResultSet.getString("udt_name")).thenReturn(udtName);
388-
when(columnsResultSet.getString("is_nullable")).thenReturn(isNullable);
390+
when(resultSet.next()).thenReturn(true, false);
391+
when(resultSet.getString("column_name")).thenReturn(colName);
392+
when(resultSet.getString("udt_name")).thenReturn(udtName);
393+
when(resultSet.getBoolean("is_nullable")).thenReturn(isNullable);
394+
when(resultSet.getBoolean("is_primary_key")).thenReturn(false);
389395
}
390396
}

0 commit comments

Comments
 (0)