Skip to content

Commit df3a484

Browse files
patelraj0602patelraj0602
committed
CODEX_AGENT at work
1 parent e869f86 commit df3a484

File tree

8 files changed

+259
-28
lines changed

8 files changed

+259
-28
lines changed

‎document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresCollection.java‎

Lines changed: 23 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,7 @@
7777
import org.hypertrace.core.documentstore.model.subdoc.SubDocumentUpdate;
7878
import org.hypertrace.core.documentstore.postgres.internal.BulkUpdateSubDocsInternalResult;
7979
import org.hypertrace.core.documentstore.postgres.model.DocumentAndId;
80+
import org.hypertrace.core.documentstore.postgres.registry.PostgresColumnRegistry;
8081
import org.hypertrace.core.documentstore.postgres.subdoc.PostgresSubDocumentUpdater;
8182
import org.hypertrace.core.documentstore.postgres.utils.PostgresUtils;
8283
import org.postgresql.util.PSQLException;
@@ -101,17 +102,33 @@ public class PostgresCollection implements Collection {
101102

102103
private final PostgresClient client;
103104
private final PostgresTableIdentifier tableIdentifier;
105+
private final PostgresColumnRegistry columnRegistry;
104106
private final PostgresSubDocumentUpdater subDocUpdater;
105107
private final PostgresQueryExecutor queryExecutor;
106108
private final UpdateValidator updateValidator;
107109

108110
public PostgresCollection(final PostgresClient client, final String collectionName) {
109-
this(client, PostgresTableIdentifier.parse(collectionName));
111+
this(client, PostgresTableIdentifier.parse(collectionName), null);
112+
}
113+
114+
public PostgresCollection(
115+
final PostgresClient client,
116+
final String collectionName,
117+
final PostgresColumnRegistry columnRegistry) {
118+
this(client, PostgresTableIdentifier.parse(collectionName), columnRegistry);
110119
}
111120

112121
PostgresCollection(final PostgresClient client, final PostgresTableIdentifier tableIdentifier) {
122+
this(client, tableIdentifier, null);
123+
}
124+
125+
PostgresCollection(
126+
final PostgresClient client,
127+
final PostgresTableIdentifier tableIdentifier,
128+
final PostgresColumnRegistry columnRegistry) {
113129
this.client = client;
114130
this.tableIdentifier = tableIdentifier;
131+
this.columnRegistry = columnRegistry;
115132
this.subDocUpdater =
116133
new PostgresSubDocumentUpdater(new PostgresQueryBuilder(this.tableIdentifier));
117134
this.queryExecutor = new PostgresQueryExecutor(this.tableIdentifier);
@@ -488,15 +505,16 @@ private CloseableIterator<Document> search(Query query, boolean removeDocumentId
488505
@Override
489506
public CloseableIterator<Document> find(
490507
final org.hypertrace.core.documentstore.query.Query query) {
491-
return queryExecutor.execute(client.getConnection(), query);
508+
return queryExecutor.execute(client.getConnection(), query, null, columnRegistry);
492509
}
493510

494511
@Override
495512
public CloseableIterator<Document> query(
496513
final org.hypertrace.core.documentstore.query.Query query, final QueryOptions queryOptions) {
497514
String flatStructureCollectionName =
498515
client.getCustomParameters().get(FLAT_STRUCTURE_COLLECTION_KEY);
499-
return queryExecutor.execute(client.getConnection(), query, flatStructureCollectionName);
516+
return queryExecutor.execute(
517+
client.getConnection(), query, flatStructureCollectionName, columnRegistry);
500518
}
501519

502520
@Override
@@ -510,7 +528,7 @@ public Optional<Document> update(
510528
try (final Connection connection = client.getPooledConnection()) {
511529
org.hypertrace.core.documentstore.postgres.query.v1.PostgresQueryParser parser =
512530
new org.hypertrace.core.documentstore.postgres.query.v1.PostgresQueryParser(
513-
tableIdentifier, query);
531+
tableIdentifier, query, columnRegistry);
514532
final String selectQuery = parser.buildSelectQueryForUpdate();
515533

516534
try (final PreparedStatement preparedStatement =
@@ -545,7 +563,7 @@ public Optional<Document> update(
545563
.build();
546564

547565
try (final CloseableIterator<Document> iterator =
548-
queryExecutor.execute(connection, findByIdQuery)) {
566+
queryExecutor.execute(connection, findByIdQuery, null, columnRegistry)) {
549567
returnDocument = getFirstDocument(iterator).orElseThrow();
550568
}
551569
} else if (updateOptions.getReturnDocumentType() == NONE) {

‎document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresDatastore.java‎

Lines changed: 82 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
import java.util.Map;
1616
import java.util.Optional;
1717
import java.util.Set;
18+
import java.util.concurrent.ConcurrentHashMap;
1819
import lombok.NonNull;
1920
import lombok.extern.slf4j.Slf4j;
2021
import org.hypertrace.core.documentstore.Collection;
@@ -24,6 +25,9 @@
2425
import org.hypertrace.core.documentstore.model.config.ConnectionConfig;
2526
import org.hypertrace.core.documentstore.model.config.DatastoreConfig;
2627
import org.hypertrace.core.documentstore.model.config.postgres.PostgresConnectionConfig;
28+
import org.hypertrace.core.documentstore.postgres.registry.PostgresColumnRegistry;
29+
import org.hypertrace.core.documentstore.postgres.registry.PostgresColumnRegistryImpl;
30+
import org.hypertrace.core.documentstore.postgres.registry.PostgresColumnType;
2731
import org.slf4j.Logger;
2832
import org.slf4j.LoggerFactory;
2933

@@ -37,6 +41,9 @@ public class PostgresDatastore implements Datastore {
3741
private final String database;
3842
private final DocStoreMetricProvider docStoreMetricProvider;
3943

44+
// Cache for PostgreSQL column registries per table
45+
private final Map<String, PostgresColumnRegistry> registryCache = new ConcurrentHashMap<>();
46+
4047
public PostgresDatastore(@NonNull final DatastoreConfig datastoreConfig) {
4148
final ConnectionConfig connectionConfig = datastoreConfig.connectionConfig();
4249

@@ -148,7 +155,81 @@ public Collection getCollection(String collectionName) {
148155
if (!tables.contains(collectionName)) {
149156
createCollection(collectionName, null);
150157
}
151-
return new PostgresCollection(client, collectionName);
158+
159+
// Create or get cached registry for this collection
160+
PostgresColumnRegistry registry = createOrGetRegistry(collectionName);
161+
162+
return new PostgresCollection(client, collectionName, registry);
163+
}
164+
165+
/**
166+
* Creates or retrieves a cached PostgresColumnRegistry for the specified collection. The registry
167+
* is cached to avoid repeated database schema queries.
168+
*
169+
* @param collectionName the collection name to create/get registry for
170+
* @return the PostgresColumnRegistry for the collection
171+
*/
172+
private PostgresColumnRegistry createOrGetRegistry(String collectionName) {
173+
return registryCache.computeIfAbsent(
174+
collectionName,
175+
tableName -> {
176+
try {
177+
PostgresColumnRegistry registry =
178+
new PostgresColumnRegistryImpl(client.getConnection(), tableName);
179+
180+
LOGGER.debug(
181+
"Created PostgresColumnRegistry for collection '{}' with {} first-class columns",
182+
tableName,
183+
registry.getAllFirstClassColumns().size());
184+
185+
return registry;
186+
} catch (SQLException e) {
187+
LOGGER.warn(
188+
"Failed to create PostgresColumnRegistry for collection '{}': {}. "
189+
+ "Falling back to JSONB-only behavior.",
190+
tableName,
191+
e.getMessage());
192+
193+
// Return an empty registry that treats all fields as JSONB
194+
return createEmptyRegistry(tableName);
195+
}
196+
});
197+
}
198+
199+
/**
200+
* Creates an empty registry that treats all fields as JSONB fields. This is used as a fallback
201+
* when registry creation fails.
202+
*
203+
* @param tableName the table name
204+
* @return an empty registry
205+
*/
206+
private PostgresColumnRegistry createEmptyRegistry(String tableName) {
207+
return new PostgresColumnRegistry() {
208+
@Override
209+
public boolean isFirstClassColumn(String fieldName) {
210+
return false; // All fields are treated as JSONB
211+
}
212+
213+
@Override
214+
public Optional<PostgresColumnType> getColumnType(String fieldName) {
215+
return Optional.empty();
216+
}
217+
218+
@Override
219+
public Optional<String> getColumnName(String fieldName) {
220+
return Optional.empty();
221+
}
222+
223+
@Override
224+
public Set<String> getAllFirstClassColumns() {
225+
return Set.of(); // No first-class columns
226+
}
227+
228+
@Override
229+
public String getTableName() {
230+
return tableName;
231+
}
232+
};
152233
}
153234

154235
@Override

‎document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresQueryExecutor.java‎

Lines changed: 24 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
import org.hypertrace.core.documentstore.postgres.PostgresCollection.PostgresResultIterator;
1414
import org.hypertrace.core.documentstore.postgres.PostgresCollection.PostgresResultIteratorWithMetaData;
1515
import org.hypertrace.core.documentstore.postgres.query.v1.transformer.PostgresQueryTransformer;
16+
import org.hypertrace.core.documentstore.postgres.registry.PostgresColumnRegistry;
1617
import org.hypertrace.core.documentstore.query.Query;
1718

1819
@Slf4j
@@ -21,22 +22,42 @@ public class PostgresQueryExecutor {
2122
private final PostgresTableIdentifier tableIdentifier;
2223

2324
public CloseableIterator<Document> execute(final Connection connection, final Query query) {
24-
return execute(connection, query, null);
25+
return execute(connection, query, null, null);
2526
}
2627

2728
public CloseableIterator<Document> execute(
2829
final Connection connection, final Query query, String flatStructureCollectionName) {
30+
return execute(connection, query, flatStructureCollectionName, null);
31+
}
32+
33+
public CloseableIterator<Document> execute(
34+
final Connection connection,
35+
final Query query,
36+
String flatStructureCollectionName,
37+
PostgresColumnRegistry columnRegistry) {
2938
final org.hypertrace.core.documentstore.postgres.query.v1.PostgresQueryParser queryParser =
3039
new org.hypertrace.core.documentstore.postgres.query.v1.PostgresQueryParser(
31-
tableIdentifier, transformAndLog(query), flatStructureCollectionName);
40+
tableIdentifier, transformAndLog(query), flatStructureCollectionName, columnRegistry);
3241
final String sqlQuery = queryParser.parse();
3342
try {
3443
final PreparedStatement preparedStatement =
3544
buildPreparedStatement(sqlQuery, queryParser.getParamsBuilder().build(), connection);
3645
log.debug("Executing executeQueryV1 sqlQuery:{}", preparedStatement.toString());
3746
final ResultSet resultSet = preparedStatement.executeQuery();
3847

39-
if ((tableIdentifier.getTableName().equals(flatStructureCollectionName))) {
48+
// Use registry-based type resolution if available
49+
boolean useBasicTypes = false;
50+
if (columnRegistry != null) {
51+
// If registry is available, use it to determine if we have first-class columns
52+
useBasicTypes = columnRegistry.hasFirstClassColumns();
53+
} else {
54+
// Fallback to flatStructureCollection for backward compatibility
55+
useBasicTypes =
56+
flatStructureCollectionName != null
57+
&& tableIdentifier.getTableName().equals(flatStructureCollectionName);
58+
}
59+
60+
if (useBasicTypes) {
4061
return new PostgresCollection.PostgresResultIteratorWithBasicTypes(resultSet);
4162
}
4263
return query.getSelections().size() > 0

‎document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/PostgresQueryParser.java‎

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import org.hypertrace.core.documentstore.postgres.query.v1.vistors.PostgresSelectTypeExpressionVisitor;
2020
import org.hypertrace.core.documentstore.postgres.query.v1.vistors.PostgresSortTypeExpressionVisitor;
2121
import org.hypertrace.core.documentstore.postgres.query.v1.vistors.PostgresUnnestFilterTypeExpressionVisitor;
22+
import org.hypertrace.core.documentstore.postgres.registry.PostgresColumnRegistry;
2223
import org.hypertrace.core.documentstore.query.Pagination;
2324
import org.hypertrace.core.documentstore.query.Query;
2425

@@ -29,6 +30,7 @@ public class PostgresQueryParser {
2930
@Getter private final PostgresTableIdentifier tableIdentifier;
3031
@Getter private final Query query;
3132
@Getter private final String flatStructureCollectionName;
33+
@Getter private final PostgresColumnRegistry columnRegistry;
3234

3335
@Setter String finalTableName;
3436
@Getter private final Builder paramsBuilder = Params.newBuilder();
@@ -47,15 +49,29 @@ public class PostgresQueryParser {
4749

4850
public PostgresQueryParser(
4951
PostgresTableIdentifier tableIdentifier, Query query, String flatStructureCollectionName) {
52+
this(tableIdentifier, query, flatStructureCollectionName, null);
53+
}
54+
55+
public PostgresQueryParser(
56+
PostgresTableIdentifier tableIdentifier,
57+
Query query,
58+
String flatStructureCollectionName,
59+
PostgresColumnRegistry columnRegistry) {
5060
this.tableIdentifier = tableIdentifier;
5161
this.query = query;
5262
this.flatStructureCollectionName = flatStructureCollectionName;
63+
this.columnRegistry = columnRegistry;
5364
this.finalTableName = tableIdentifier.toString();
5465
toPgColumnTransformer = new FieldToPgColumnTransformer(this);
5566
}
5667

5768
public PostgresQueryParser(PostgresTableIdentifier tableIdentifier, Query query) {
58-
this(tableIdentifier, query, null);
69+
this(tableIdentifier, query, null, null);
70+
}
71+
72+
public PostgresQueryParser(
73+
PostgresTableIdentifier tableIdentifier, Query query, PostgresColumnRegistry columnRegistry) {
74+
this(tableIdentifier, query, null, columnRegistry);
5975
}
6076

6177
public String parse() {

‎document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/parser/filter/PostgresNotContainsRelationalFilterParser.java‎

Lines changed: 29 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,9 @@
11
package org.hypertrace.core.documentstore.postgres.query.v1.parser.filter;
22

3+
import org.hypertrace.core.documentstore.expression.impl.IdentifierExpression;
34
import org.hypertrace.core.documentstore.expression.impl.RelationalExpression;
45
import org.hypertrace.core.documentstore.postgres.query.v1.parser.filter.nonjson.field.PostgresContainsRelationalFilterParserNonJsonField;
6+
import org.hypertrace.core.documentstore.postgres.registry.PostgresColumnRegistry;
57

68
class PostgresNotContainsRelationalFilterParser implements PostgresRelationalFilterParser {
79
private static final PostgresContainsRelationalFilterParser jsonContainsParser =
@@ -14,10 +16,9 @@ public String parse(
1416
final RelationalExpression expression, final PostgresRelationalFilterContext context) {
1517
final String parsedLhs = expression.getLhs().accept(context.lhsParser());
1618

17-
String flatStructureCollection = context.getFlatStructureCollectionName();
18-
boolean isFirstClassField =
19-
flatStructureCollection != null
20-
&& flatStructureCollection.equals(context.getTableIdentifier().getTableName());
19+
// Extract field name and determine if it's a first-class field
20+
String fieldName = extractFieldName(expression);
21+
boolean isFirstClassField = determineIfFirstClassField(fieldName, context);
2122

2223
if (isFirstClassField) {
2324
// Use the non-JSON logic for first-class fields
@@ -29,4 +30,28 @@ public String parse(
2930
return String.format("%s IS NULL OR NOT %s @> ?::jsonb", parsedLhs, parsedLhs);
3031
}
3132
}
33+
34+
/** Extracts the field name from the left-hand side of a RelationalExpression. */
35+
private String extractFieldName(RelationalExpression expression) {
36+
if (expression.getLhs() instanceof IdentifierExpression) {
37+
return ((IdentifierExpression) expression.getLhs()).getName();
38+
}
39+
return null;
40+
}
41+
42+
/** Determines if a field is a first-class column using registry-based lookup with fallback. */
43+
private boolean determineIfFirstClassField(
44+
String fieldName, PostgresRelationalFilterContext context) {
45+
PostgresColumnRegistry registry = context.getColumnRegistry();
46+
47+
// Use registry-based type resolution if available
48+
if (registry != null && fieldName != null) {
49+
return registry.isFirstClassColumn(fieldName);
50+
} else {
51+
// Fallback to flatStructureCollection for backward compatibility
52+
String flatStructureCollection = context.getFlatStructureCollectionName();
53+
return flatStructureCollection != null
54+
&& flatStructureCollection.equals(context.getTableIdentifier().getTableName());
55+
}
56+
}
3257
}

‎document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/parser/filter/PostgresNotInRelationalFilterParser.java‎

Lines changed: 33 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import org.hypertrace.core.documentstore.expression.impl.RelationalExpression;
44
import org.hypertrace.core.documentstore.postgres.query.v1.parser.filter.nonjson.field.PostgresInRelationalFilterParserNonJsonField;
5+
import org.hypertrace.core.documentstore.postgres.registry.PostgresColumnRegistry;
56

67
class PostgresNotInRelationalFilterParser implements PostgresRelationalFilterParser {
78
private static final PostgresInRelationalFilterParserInterface jsonFieldInFilterParser =
@@ -22,11 +23,39 @@ public String parse(
2223

2324
private PostgresInRelationalFilterParserInterface getInFilterParser(
2425
PostgresRelationalFilterContext context) {
25-
String flatStructureCollection = context.getFlatStructureCollectionName();
26-
boolean isFirstClassField =
27-
flatStructureCollection != null
28-
&& flatStructureCollection.equals(context.getTableIdentifier().getTableName());
26+
// Extract field name from the expression
27+
String fieldName = extractFieldName(context);
28+
boolean isFirstClassField = determineIfFirstClassField(fieldName, context);
2929

3030
return isFirstClassField ? nonJsonFieldInFilterParser : jsonFieldInFilterParser;
3131
}
32+
33+
/**
34+
* Extracts the field name from the context's current expression. This is a simplified approach -
35+
* in a full implementation, we'd need access to the expression.
36+
*/
37+
private String extractFieldName(PostgresRelationalFilterContext context) {
38+
// Note: This is a limitation of the current design - we don't have direct access to the
39+
// expression here.
40+
// For now, we'll return null and rely on the fallback logic.
41+
// In a full refactor, we'd pass the field name through the context or restructure the parser
42+
// hierarchy.
43+
return null;
44+
}
45+
46+
/** Determines if a field is a first-class column using registry-based lookup with fallback. */
47+
private boolean determineIfFirstClassField(
48+
String fieldName, PostgresRelationalFilterContext context) {
49+
PostgresColumnRegistry registry = context.getColumnRegistry();
50+
51+
// Use registry-based type resolution if available
52+
if (registry != null && fieldName != null) {
53+
return registry.isFirstClassColumn(fieldName);
54+
} else {
55+
// Fallback to flatStructureCollection for backward compatibility
56+
String flatStructureCollection = context.getFlatStructureCollectionName();
57+
return flatStructureCollection != null
58+
&& flatStructureCollection.equals(context.getTableIdentifier().getTableName());
59+
}
60+
}
3261
}

0 commit comments

Comments
 (0)