From 20d611663caea5b8dc69771ea51d0e3b3c8eb925 Mon Sep 17 00:00:00 2001 From: Patel Raj Pareshkumar Date: Thu, 14 Aug 2025 16:12:51 +0530 Subject: [PATCH] fixed integration test --- .../postgres/PostgresCollection.java | 59 ++++++- .../postgres/PostgresDatastore.java | 83 +++++++++- .../postgres/PostgresQueryBuilder.java | 5 +- .../postgres/PostgresQueryExecutor.java | 27 ++- .../query/v1/PostgresQueryParser.java | 18 +- ...gresNotContainsRelationalFilterParser.java | 33 +++- .../PostgresNotInRelationalFilterParser.java | 37 ++++- ...gresRelationalFilterParserFactoryImpl.java | 39 ++++- .../FieldToPgColumnTransformer.java | 22 ++- .../postgres/registry/PostgresColumnInfo.java | 75 +++++++++ .../registry/PostgresColumnRegistry.java | 66 ++++++++ .../registry/PostgresColumnRegistryImpl.java | 156 ++++++++++++++++++ .../postgres/registry/PostgresColumnType.java | 88 ++++++++++ 13 files changed, 673 insertions(+), 35 deletions(-) create mode 100644 document-store/src/main/java/org/hypertrace/core/documentstore/postgres/registry/PostgresColumnInfo.java create mode 100644 document-store/src/main/java/org/hypertrace/core/documentstore/postgres/registry/PostgresColumnRegistry.java create mode 100644 document-store/src/main/java/org/hypertrace/core/documentstore/postgres/registry/PostgresColumnRegistryImpl.java create mode 100644 document-store/src/main/java/org/hypertrace/core/documentstore/postgres/registry/PostgresColumnType.java diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresCollection.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresCollection.java index 5582889a..dfe5c608 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresCollection.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresCollection.java @@ -77,6 +77,7 @@ import org.hypertrace.core.documentstore.model.subdoc.SubDocumentUpdate; import org.hypertrace.core.documentstore.postgres.internal.BulkUpdateSubDocsInternalResult; import org.hypertrace.core.documentstore.postgres.model.DocumentAndId; +import org.hypertrace.core.documentstore.postgres.registry.PostgresColumnRegistry; import org.hypertrace.core.documentstore.postgres.subdoc.PostgresSubDocumentUpdater; import org.hypertrace.core.documentstore.postgres.utils.PostgresUtils; import org.postgresql.util.PSQLException; @@ -101,19 +102,36 @@ public class PostgresCollection implements Collection { private final PostgresClient client; private final PostgresTableIdentifier tableIdentifier; + private final PostgresColumnRegistry columnRegistry; private final PostgresSubDocumentUpdater subDocUpdater; private final PostgresQueryExecutor queryExecutor; private final UpdateValidator updateValidator; public PostgresCollection(final PostgresClient client, final String collectionName) { - this(client, PostgresTableIdentifier.parse(collectionName)); + this(client, PostgresTableIdentifier.parse(collectionName), null); + } + + public PostgresCollection( + final PostgresClient client, + final String collectionName, + final PostgresColumnRegistry columnRegistry) { + this(client, PostgresTableIdentifier.parse(collectionName), columnRegistry); } PostgresCollection(final PostgresClient client, final PostgresTableIdentifier tableIdentifier) { + this(client, tableIdentifier, null); + } + + PostgresCollection( + final PostgresClient client, + final PostgresTableIdentifier tableIdentifier, + final PostgresColumnRegistry columnRegistry) { this.client = client; this.tableIdentifier = tableIdentifier; + this.columnRegistry = columnRegistry; this.subDocUpdater = - new PostgresSubDocumentUpdater(new PostgresQueryBuilder(this.tableIdentifier)); + new PostgresSubDocumentUpdater( + new PostgresQueryBuilder(this.tableIdentifier, this.columnRegistry)); this.queryExecutor = new PostgresQueryExecutor(this.tableIdentifier); this.updateValidator = new CommonUpdateValidator(); } @@ -488,7 +506,7 @@ private CloseableIterator search(Query query, boolean removeDocumentId @Override public CloseableIterator find( final org.hypertrace.core.documentstore.query.Query query) { - return queryExecutor.execute(client.getConnection(), query); + return queryExecutor.execute(client.getConnection(), query, null, columnRegistry); } @Override @@ -496,7 +514,8 @@ public CloseableIterator query( final org.hypertrace.core.documentstore.query.Query query, final QueryOptions queryOptions) { String flatStructureCollectionName = client.getCustomParameters().get(FLAT_STRUCTURE_COLLECTION_KEY); - return queryExecutor.execute(client.getConnection(), query, flatStructureCollectionName); + return queryExecutor.execute( + client.getConnection(), query, flatStructureCollectionName, columnRegistry); } @Override @@ -510,7 +529,7 @@ public Optional update( try (final Connection connection = client.getPooledConnection()) { org.hypertrace.core.documentstore.postgres.query.v1.PostgresQueryParser parser = new org.hypertrace.core.documentstore.postgres.query.v1.PostgresQueryParser( - tableIdentifier, query); + tableIdentifier, query, columnRegistry); final String selectQuery = parser.buildSelectQueryForUpdate(); try (final PreparedStatement preparedStatement = @@ -545,7 +564,7 @@ public Optional update( .build(); try (final CloseableIterator iterator = - queryExecutor.execute(connection, findByIdQuery)) { + queryExecutor.execute(connection, findByIdQuery, null, columnRegistry)) { returnDocument = getFirstDocument(iterator).orElseThrow(); } } else if (updateOptions.getReturnDocumentType() == NONE) { @@ -1568,4 +1587,32 @@ private static Optional mapValueToJsonNode(int columnType, String colu return Optional.empty(); } } + + /** + * Iterator that extracts clean documents from the 'document' column for queries without + * selections, removing the wrapper structure. + */ + static class PostgresResultIteratorCleanDocument extends PostgresResultIterator { + + public PostgresResultIteratorCleanDocument(ResultSet resultSet) { + super(resultSet, true); + } + + @Override + protected Document prepareDocument() throws SQLException, IOException { + // For queries without selections, extract the document content directly + String documentString = resultSet.getString(DOCUMENT); + if (documentString != null) { + ObjectNode jsonNode = (ObjectNode) MAPPER.readTree(documentString); + // Remove document ID if needed + if (shouldRemoveDocumentId()) { + jsonNode.remove(DOCUMENT_ID); + } + return new JSONDocument(MAPPER.writeValueAsString(jsonNode)); + } + + // Fallback to empty document if no document column + return new JSONDocument("{}"); + } + } } diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresDatastore.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresDatastore.java index 84678ca7..22479760 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresDatastore.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresDatastore.java @@ -15,6 +15,7 @@ import java.util.Map; import java.util.Optional; import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; import lombok.NonNull; import lombok.extern.slf4j.Slf4j; import org.hypertrace.core.documentstore.Collection; @@ -24,6 +25,9 @@ import org.hypertrace.core.documentstore.model.config.ConnectionConfig; import org.hypertrace.core.documentstore.model.config.DatastoreConfig; import org.hypertrace.core.documentstore.model.config.postgres.PostgresConnectionConfig; +import org.hypertrace.core.documentstore.postgres.registry.PostgresColumnRegistry; +import org.hypertrace.core.documentstore.postgres.registry.PostgresColumnRegistryImpl; +import org.hypertrace.core.documentstore.postgres.registry.PostgresColumnType; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -37,6 +41,9 @@ public class PostgresDatastore implements Datastore { private final String database; private final DocStoreMetricProvider docStoreMetricProvider; + // Cache for PostgreSQL column registries per table + private final Map registryCache = new ConcurrentHashMap<>(); + public PostgresDatastore(@NonNull final DatastoreConfig datastoreConfig) { final ConnectionConfig connectionConfig = datastoreConfig.connectionConfig(); @@ -148,7 +155,81 @@ public Collection getCollection(String collectionName) { if (!tables.contains(collectionName)) { createCollection(collectionName, null); } - return new PostgresCollection(client, collectionName); + + // Create or get cached registry for this collection + PostgresColumnRegistry registry = createOrGetRegistry(collectionName); + + return new PostgresCollection(client, collectionName, registry); + } + + /** + * Creates or retrieves a cached PostgresColumnRegistry for the specified collection. The registry + * is cached to avoid repeated database schema queries. + * + * @param collectionName the collection name to create/get registry for + * @return the PostgresColumnRegistry for the collection + */ + private PostgresColumnRegistry createOrGetRegistry(String collectionName) { + return registryCache.computeIfAbsent( + collectionName, + tableName -> { + try { + PostgresColumnRegistry registry = + new PostgresColumnRegistryImpl(client.getConnection(), tableName); + + LOGGER.debug( + "Created PostgresColumnRegistry for collection '{}' with {} first-class columns", + tableName, + registry.getAllFirstClassColumns().size()); + + return registry; + } catch (SQLException e) { + LOGGER.warn( + "Failed to create PostgresColumnRegistry for collection '{}': {}. " + + "Falling back to JSONB-only behavior.", + tableName, + e.getMessage()); + + // Return an empty registry that treats all fields as JSONB + return createEmptyRegistry(tableName); + } + }); + } + + /** + * Creates an empty registry that treats all fields as JSONB fields. This is used as a fallback + * when registry creation fails. + * + * @param tableName the table name + * @return an empty registry + */ + private PostgresColumnRegistry createEmptyRegistry(String tableName) { + return new PostgresColumnRegistry() { + @Override + public boolean isFirstClassColumn(String fieldName) { + return false; // All fields are treated as JSONB + } + + @Override + public Optional getColumnType(String fieldName) { + return Optional.empty(); + } + + @Override + public Optional getColumnName(String fieldName) { + return Optional.empty(); + } + + @Override + public Set getAllFirstClassColumns() { + return Set.of(); // No first-class columns + } + + @Override + public String getTableName() { + return tableName; + } + }; } @Override diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresQueryBuilder.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresQueryBuilder.java index 9edb5ab3..b1460348 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresQueryBuilder.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresQueryBuilder.java @@ -21,6 +21,7 @@ import org.hypertrace.core.documentstore.model.subdoc.UpdateOperator; import org.hypertrace.core.documentstore.postgres.Params.Builder; import org.hypertrace.core.documentstore.postgres.query.v1.PostgresQueryParser; +import org.hypertrace.core.documentstore.postgres.registry.PostgresColumnRegistry; import org.hypertrace.core.documentstore.postgres.update.parser.PostgresAddToListIfAbsentParser; import org.hypertrace.core.documentstore.postgres.update.parser.PostgresAddValueParser; import org.hypertrace.core.documentstore.postgres.update.parser.PostgresAppendToListParser; @@ -44,12 +45,14 @@ public class PostgresQueryBuilder { entry(APPEND_TO_LIST, new PostgresAppendToListParser())); @Getter private final PostgresTableIdentifier tableIdentifier; + private final PostgresColumnRegistry columnRegistry; public String getSubDocUpdateQuery( final Query query, final Collection updates, final Params.Builder paramBuilder) { - final PostgresQueryParser baseQueryParser = new PostgresQueryParser(tableIdentifier, query); + final PostgresQueryParser baseQueryParser = + new PostgresQueryParser(tableIdentifier, query, columnRegistry); String selectQuery = String.format( "(SELECT %s, %s FROM %s AS t0 %s)", diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresQueryExecutor.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresQueryExecutor.java index 389a52b8..3baf1814 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresQueryExecutor.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresQueryExecutor.java @@ -10,9 +10,9 @@ import lombok.extern.slf4j.Slf4j; import org.hypertrace.core.documentstore.CloseableIterator; import org.hypertrace.core.documentstore.Document; -import org.hypertrace.core.documentstore.postgres.PostgresCollection.PostgresResultIterator; import org.hypertrace.core.documentstore.postgres.PostgresCollection.PostgresResultIteratorWithMetaData; import org.hypertrace.core.documentstore.postgres.query.v1.transformer.PostgresQueryTransformer; +import org.hypertrace.core.documentstore.postgres.registry.PostgresColumnRegistry; import org.hypertrace.core.documentstore.query.Query; @Slf4j @@ -21,14 +21,22 @@ public class PostgresQueryExecutor { private final PostgresTableIdentifier tableIdentifier; public CloseableIterator execute(final Connection connection, final Query query) { - return execute(connection, query, null); + return execute(connection, query, null, null); } public CloseableIterator execute( final Connection connection, final Query query, String flatStructureCollectionName) { + return execute(connection, query, flatStructureCollectionName, null); + } + + public CloseableIterator execute( + final Connection connection, + final Query query, + String flatStructureCollectionName, + PostgresColumnRegistry columnRegistry) { final org.hypertrace.core.documentstore.postgres.query.v1.PostgresQueryParser queryParser = new org.hypertrace.core.documentstore.postgres.query.v1.PostgresQueryParser( - tableIdentifier, transformAndLog(query), flatStructureCollectionName); + tableIdentifier, transformAndLog(query), flatStructureCollectionName, columnRegistry); final String sqlQuery = queryParser.parse(); try { final PreparedStatement preparedStatement = @@ -36,12 +44,15 @@ public CloseableIterator execute( log.debug("Executing executeQueryV1 sqlQuery:{}", preparedStatement.toString()); final ResultSet resultSet = preparedStatement.executeQuery(); - if ((tableIdentifier.getTableName().equals(flatStructureCollectionName))) { - return new PostgresCollection.PostgresResultIteratorWithBasicTypes(resultSet); + // For queries with selections, use PostgresResultIteratorWithMetaData + // as it properly handles nested field decoding + if (query.getSelections().size() > 0) { + return new PostgresResultIteratorWithMetaData(resultSet); } - return query.getSelections().size() > 0 - ? new PostgresResultIteratorWithMetaData(resultSet) - : new PostgresResultIterator(resultSet); + + // For queries without selections, use a custom iterator that extracts clean documents + // from the 'document' column instead of returning wrapped results + return new PostgresCollection.PostgresResultIteratorCleanDocument(resultSet); } catch (SQLException e) { log.error( "SQLException querying documents. original query: " + query + ", sql query:" + sqlQuery, diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/PostgresQueryParser.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/PostgresQueryParser.java index c9a7a55b..53271d6c 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/PostgresQueryParser.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/PostgresQueryParser.java @@ -19,6 +19,7 @@ import org.hypertrace.core.documentstore.postgres.query.v1.vistors.PostgresSelectTypeExpressionVisitor; import org.hypertrace.core.documentstore.postgres.query.v1.vistors.PostgresSortTypeExpressionVisitor; import org.hypertrace.core.documentstore.postgres.query.v1.vistors.PostgresUnnestFilterTypeExpressionVisitor; +import org.hypertrace.core.documentstore.postgres.registry.PostgresColumnRegistry; import org.hypertrace.core.documentstore.query.Pagination; import org.hypertrace.core.documentstore.query.Query; @@ -29,6 +30,7 @@ public class PostgresQueryParser { @Getter private final PostgresTableIdentifier tableIdentifier; @Getter private final Query query; @Getter private final String flatStructureCollectionName; + @Getter private final PostgresColumnRegistry columnRegistry; @Setter String finalTableName; @Getter private final Builder paramsBuilder = Params.newBuilder(); @@ -47,15 +49,29 @@ public class PostgresQueryParser { public PostgresQueryParser( PostgresTableIdentifier tableIdentifier, Query query, String flatStructureCollectionName) { + this(tableIdentifier, query, flatStructureCollectionName, null); + } + + public PostgresQueryParser( + PostgresTableIdentifier tableIdentifier, + Query query, + String flatStructureCollectionName, + PostgresColumnRegistry columnRegistry) { this.tableIdentifier = tableIdentifier; this.query = query; this.flatStructureCollectionName = flatStructureCollectionName; + this.columnRegistry = columnRegistry; this.finalTableName = tableIdentifier.toString(); toPgColumnTransformer = new FieldToPgColumnTransformer(this); } public PostgresQueryParser(PostgresTableIdentifier tableIdentifier, Query query) { - this(tableIdentifier, query, null); + this(tableIdentifier, query, null, null); + } + + public PostgresQueryParser( + PostgresTableIdentifier tableIdentifier, Query query, PostgresColumnRegistry columnRegistry) { + this(tableIdentifier, query, null, columnRegistry); } public String parse() { diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/parser/filter/PostgresNotContainsRelationalFilterParser.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/parser/filter/PostgresNotContainsRelationalFilterParser.java index 15eea744..5ee6f2c8 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/parser/filter/PostgresNotContainsRelationalFilterParser.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/parser/filter/PostgresNotContainsRelationalFilterParser.java @@ -1,7 +1,9 @@ package org.hypertrace.core.documentstore.postgres.query.v1.parser.filter; +import org.hypertrace.core.documentstore.expression.impl.IdentifierExpression; import org.hypertrace.core.documentstore.expression.impl.RelationalExpression; import org.hypertrace.core.documentstore.postgres.query.v1.parser.filter.nonjson.field.PostgresContainsRelationalFilterParserNonJsonField; +import org.hypertrace.core.documentstore.postgres.registry.PostgresColumnRegistry; class PostgresNotContainsRelationalFilterParser implements PostgresRelationalFilterParser { private static final PostgresContainsRelationalFilterParser jsonContainsParser = @@ -14,10 +16,9 @@ public String parse( final RelationalExpression expression, final PostgresRelationalFilterContext context) { final String parsedLhs = expression.getLhs().accept(context.lhsParser()); - String flatStructureCollection = context.getFlatStructureCollectionName(); - boolean isFirstClassField = - flatStructureCollection != null - && flatStructureCollection.equals(context.getTableIdentifier().getTableName()); + // Extract field name and determine if it's a first-class field + String fieldName = extractFieldName(expression); + boolean isFirstClassField = determineIfFirstClassField(fieldName, context); if (isFirstClassField) { // Use the non-JSON logic for first-class fields @@ -29,4 +30,28 @@ public String parse( return String.format("%s IS NULL OR NOT %s @> ?::jsonb", parsedLhs, parsedLhs); } } + + /** Extracts the field name from the left-hand side of a RelationalExpression. */ + private String extractFieldName(RelationalExpression expression) { + if (expression.getLhs() instanceof IdentifierExpression) { + return ((IdentifierExpression) expression.getLhs()).getName(); + } + return null; + } + + /** Determines if a field is a first-class column using registry-based lookup with fallback. */ + private boolean determineIfFirstClassField( + String fieldName, PostgresRelationalFilterContext context) { + PostgresColumnRegistry registry = context.getColumnRegistry(); + + // Use registry-based type resolution if available + if (registry != null && fieldName != null) { + return registry.isFirstClassColumn(fieldName); + } else { + // Fallback to flatStructureCollection for backward compatibility + String flatStructureCollection = context.getFlatStructureCollectionName(); + return flatStructureCollection != null + && flatStructureCollection.equals(context.getTableIdentifier().getTableName()); + } + } } diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/parser/filter/PostgresNotInRelationalFilterParser.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/parser/filter/PostgresNotInRelationalFilterParser.java index 92b2759c..ea691e63 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/parser/filter/PostgresNotInRelationalFilterParser.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/parser/filter/PostgresNotInRelationalFilterParser.java @@ -2,6 +2,7 @@ import org.hypertrace.core.documentstore.expression.impl.RelationalExpression; import org.hypertrace.core.documentstore.postgres.query.v1.parser.filter.nonjson.field.PostgresInRelationalFilterParserNonJsonField; +import org.hypertrace.core.documentstore.postgres.registry.PostgresColumnRegistry; class PostgresNotInRelationalFilterParser implements PostgresRelationalFilterParser { private static final PostgresInRelationalFilterParserInterface jsonFieldInFilterParser = @@ -22,11 +23,39 @@ public String parse( private PostgresInRelationalFilterParserInterface getInFilterParser( PostgresRelationalFilterContext context) { - String flatStructureCollection = context.getFlatStructureCollectionName(); - boolean isFirstClassField = - flatStructureCollection != null - && flatStructureCollection.equals(context.getTableIdentifier().getTableName()); + // Extract field name from the expression + String fieldName = extractFieldName(context); + boolean isFirstClassField = determineIfFirstClassField(fieldName, context); return isFirstClassField ? nonJsonFieldInFilterParser : jsonFieldInFilterParser; } + + /** + * Extracts the field name from the context's current expression. This is a simplified approach - + * in a full implementation, we'd need access to the expression. + */ + private String extractFieldName(PostgresRelationalFilterContext context) { + // Note: This is a limitation of the current design - we don't have direct access to the + // expression here. + // For now, we'll return null and rely on the fallback logic. + // In a full refactor, we'd pass the field name through the context or restructure the parser + // hierarchy. + return null; + } + + /** Determines if a field is a first-class column using registry-based lookup with fallback. */ + private boolean determineIfFirstClassField( + String fieldName, PostgresRelationalFilterContext context) { + PostgresColumnRegistry registry = context.getColumnRegistry(); + + // Use registry-based type resolution if available + if (registry != null && fieldName != null) { + return registry.isFirstClassColumn(fieldName); + } else { + // Fallback to flatStructureCollection for backward compatibility + String flatStructureCollection = context.getFlatStructureCollectionName(); + return flatStructureCollection != null + && flatStructureCollection.equals(context.getTableIdentifier().getTableName()); + } + } } diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/parser/filter/PostgresRelationalFilterParserFactoryImpl.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/parser/filter/PostgresRelationalFilterParserFactoryImpl.java index 8b4ef735..5a50d85d 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/parser/filter/PostgresRelationalFilterParserFactoryImpl.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/parser/filter/PostgresRelationalFilterParserFactoryImpl.java @@ -12,11 +12,13 @@ import com.google.common.collect.Maps; import java.util.Map; +import org.hypertrace.core.documentstore.expression.impl.IdentifierExpression; import org.hypertrace.core.documentstore.expression.impl.RelationalExpression; import org.hypertrace.core.documentstore.expression.operators.RelationalOperator; import org.hypertrace.core.documentstore.postgres.query.v1.PostgresQueryParser; import org.hypertrace.core.documentstore.postgres.query.v1.parser.filter.nonjson.field.PostgresContainsRelationalFilterParserNonJsonField; import org.hypertrace.core.documentstore.postgres.query.v1.parser.filter.nonjson.field.PostgresInRelationalFilterParserNonJsonField; +import org.hypertrace.core.documentstore.postgres.registry.PostgresColumnRegistry; public class PostgresRelationalFilterParserFactoryImpl implements PostgresRelationalFilterParserFactory { @@ -51,11 +53,9 @@ public class PostgresRelationalFilterParserFactoryImpl public PostgresRelationalFilterParser parser( final RelationalExpression expression, final PostgresQueryParser postgresQueryParser) { - String flatStructureCollection = postgresQueryParser.getFlatStructureCollectionName(); - boolean isFirstClassField = - flatStructureCollection != null - && flatStructureCollection.equals( - postgresQueryParser.getTableIdentifier().getTableName()); + // Extract field name from the left-hand side of the expression + String fieldName = extractFieldName(expression); + boolean isFirstClassField = determineIfFirstClassField(fieldName, postgresQueryParser); if (expression.getOperator() == CONTAINS) { return isFirstClassField ? nonJsonFieldContainsParser : jsonFieldContainsParser; @@ -65,4 +65,33 @@ public PostgresRelationalFilterParser parser( return parserMap.getOrDefault(expression.getOperator(), postgresStandardRelationalFilterParser); } + + /** + * Extracts the field name from the left-hand side of a RelationalExpression. Currently supports + * IdentifierExpression; returns null for other types. + */ + private String extractFieldName(RelationalExpression expression) { + if (expression.getLhs() instanceof IdentifierExpression) { + return ((IdentifierExpression) expression.getLhs()).getName(); + } + // For other expression types (functions, etc.), we can't easily extract a field name + return null; + } + + /** Determines if a field is a first-class column using registry-based lookup with fallback. */ + private boolean determineIfFirstClassField( + String fieldName, PostgresQueryParser postgresQueryParser) { + PostgresColumnRegistry registry = postgresQueryParser.getColumnRegistry(); + + // Use registry-based type resolution if available + if (registry != null && fieldName != null) { + return registry.isFirstClassColumn(fieldName); + } else { + // Fallback to flatStructureCollection for backward compatibility + String flatStructureCollection = postgresQueryParser.getFlatStructureCollectionName(); + return flatStructureCollection != null + && flatStructureCollection.equals( + postgresQueryParser.getTableIdentifier().getTableName()); + } + } } diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/transformer/FieldToPgColumnTransformer.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/transformer/FieldToPgColumnTransformer.java index 7e5769ef..d3c7ccba 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/transformer/FieldToPgColumnTransformer.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/transformer/FieldToPgColumnTransformer.java @@ -4,6 +4,7 @@ import java.util.Optional; import org.apache.commons.lang3.StringUtils; import org.hypertrace.core.documentstore.postgres.query.v1.PostgresQueryParser; +import org.hypertrace.core.documentstore.postgres.registry.PostgresColumnRegistry; import org.hypertrace.core.documentstore.postgres.utils.PostgresUtils; public class FieldToPgColumnTransformer { @@ -16,11 +17,22 @@ public FieldToPgColumnTransformer(PostgresQueryParser postgresQueryParser) { } public FieldToPgColumn transform(String orgFieldName) { - // TODO: Forcing to map to the first class fields - String flatStructureCollection = postgresQueryParser.getFlatStructureCollectionName(); - if (flatStructureCollection != null - && flatStructureCollection.equals( - postgresQueryParser.getTableIdentifier().getTableName())) { + // Use registry-based type resolution if available + PostgresColumnRegistry registry = postgresQueryParser.getColumnRegistry(); + boolean isFirstClassField = false; + + if (registry != null) { + isFirstClassField = registry.isFirstClassColumn(orgFieldName); + } else { + // Fallback to flatStructureCollection for backward compatibility + String flatStructureCollection = postgresQueryParser.getFlatStructureCollectionName(); + isFirstClassField = + flatStructureCollection != null + && flatStructureCollection.equals( + postgresQueryParser.getTableIdentifier().getTableName()); + } + + if (isFirstClassField) { return new FieldToPgColumn(null, PostgresUtils.wrapFieldNamesWithDoubleQuotes(orgFieldName)); } Optional parentField = diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/registry/PostgresColumnInfo.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/registry/PostgresColumnInfo.java new file mode 100644 index 00000000..95e4a053 --- /dev/null +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/registry/PostgresColumnInfo.java @@ -0,0 +1,75 @@ +package org.hypertrace.core.documentstore.postgres.registry; + +import java.util.Objects; + +/** + * Immutable data class representing metadata about a PostgreSQL column. Contains the column name + * and its corresponding PostgreSQL data type. + */ +public final class PostgresColumnInfo { + + private final String columnName; + private final PostgresColumnType columnType; + + /** + * Creates a new PostgresColumnInfo instance. + * + * @param columnName the PostgreSQL column name (must not be null) + * @param columnType the PostgreSQL column type (must not be null) + * @throws IllegalArgumentException if columnName or columnType is null + */ + public PostgresColumnInfo(String columnName, PostgresColumnType columnType) { + this.columnName = Objects.requireNonNull(columnName, "Column name cannot be null"); + this.columnType = Objects.requireNonNull(columnType, "Column type cannot be null"); + } + + /** + * Gets the PostgreSQL column name. + * + * @return the column name + */ + public String getColumnName() { + return columnName; + } + + /** + * Gets the PostgreSQL column type. + * + * @return the column type + */ + public PostgresColumnType getColumnType() { + return columnType; + } + + /** + * Checks if this column represents a first-class field (non-JSONB). + * + * @return true if this is a first-class field, false if it's a JSONB field + */ + public boolean isFirstClassField() { + return columnType.isFirstClassField(); + } + + @Override + public boolean equals(Object obj) { + if (this == obj) { + return true; + } + if (obj == null || getClass() != obj.getClass()) { + return false; + } + PostgresColumnInfo that = (PostgresColumnInfo) obj; + return Objects.equals(columnName, that.columnName) && columnType == that.columnType; + } + + @Override + public int hashCode() { + return Objects.hash(columnName, columnType); + } + + @Override + public String toString() { + return String.format( + "PostgresColumnInfo{columnName='%s', columnType=%s}", columnName, columnType); + } +} diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/registry/PostgresColumnRegistry.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/registry/PostgresColumnRegistry.java new file mode 100644 index 00000000..764d516b --- /dev/null +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/registry/PostgresColumnRegistry.java @@ -0,0 +1,66 @@ +package org.hypertrace.core.documentstore.postgres.registry; + +import java.util.Optional; +import java.util.Set; + +/** + * Registry interface for PostgreSQL column type information. Provides metadata about table columns + * to determine appropriate query generation strategies. + * + *

This registry replaces hardcoded column lists and enables dynamic, database-driven decisions + * about whether to use JSON-based or native PostgreSQL column access. + */ +public interface PostgresColumnRegistry { + + /** + * Determines if the specified field name corresponds to a first-class column (i.e., a native + * PostgreSQL column rather than a field within a JSONB document). + * + * @param fieldName the field name to check + * @return true if this is a first-class column, false if it should be accessed via JSONB + */ + boolean isFirstClassColumn(String fieldName); + + /** + * Gets the PostgreSQL column type for the specified field name. + * + * @param fieldName the field name to look up + * @return the PostgreSQL column type, or empty if the field is not a first-class column + */ + Optional getColumnType(String fieldName); + + /** + * Gets the actual PostgreSQL column name for the specified field name. In most cases, this will + * be the same as the field name, but this method allows for potential field name transformations + * or aliases. + * + * @param fieldName the field name to look up + * @return the PostgreSQL column name, or empty if the field is not a first-class column + */ + Optional getColumnName(String fieldName); + + /** + * Gets all first-class column names in this table. This is useful for debugging and validation + * purposes. + * + * @return a set of all first-class column names + */ + Set getAllFirstClassColumns(); + + /** + * Gets the table name this registry represents. + * + * @return the table name + */ + String getTableName(); + + /** + * Checks if this registry has any first-class columns. If false, all fields should be accessed + * via JSONB. + * + * @return true if there are first-class columns, false otherwise + */ + default boolean hasFirstClassColumns() { + return !getAllFirstClassColumns().isEmpty(); + } +} diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/registry/PostgresColumnRegistryImpl.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/registry/PostgresColumnRegistryImpl.java new file mode 100644 index 00000000..50905c56 --- /dev/null +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/registry/PostgresColumnRegistryImpl.java @@ -0,0 +1,156 @@ +package org.hypertrace.core.documentstore.postgres.registry; + +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import lombok.extern.slf4j.Slf4j; +import org.hypertrace.core.documentstore.postgres.PostgresTableIdentifier; + +/** + * Implementation of PostgresColumnRegistry that queries the database schema to build column type + * mappings dynamically. + * + *

This implementation queries the information_schema.columns table to discover the actual column + * types in the PostgreSQL table and maps them to the appropriate PostgresColumnType enum values. + */ +@Slf4j +public class PostgresColumnRegistryImpl implements PostgresColumnRegistry { + + private final String tableName; + private final Map columnMappings; + + /** + * Creates a new PostgresColumnRegistryImpl by querying the database schema. + * + * @param connection the database connection to use for schema queries + * @param tableIdentifier the table identifier to query schema for + * @throws SQLException if there's an error querying the database schema + */ + public PostgresColumnRegistryImpl(Connection connection, PostgresTableIdentifier tableIdentifier) + throws SQLException { + this.tableName = tableIdentifier.getTableName(); + this.columnMappings = buildColumnMappings(connection, tableIdentifier); + System.out.println("-- PostgresColumnRegistryImpl --"); + System.out.printf("Table Name: %s\n", this.tableName); + System.out.println(this.columnMappings); + System.out.println("-- end --"); + System.out.println(); + + log.debug( + "Created PostgresColumnRegistry for table '{}' with {} first-class columns: {}", + tableName, + columnMappings.size(), + columnMappings.keySet()); + } + + /** + * Creates a new PostgresColumnRegistryImpl by querying the database schema. + * + * @param connection the database connection to use for schema queries + * @param collectionName the collection name (table name) to query schema for + * @throws SQLException if there's an error querying the database schema + */ + public PostgresColumnRegistryImpl(Connection connection, String collectionName) + throws SQLException { + this(connection, PostgresTableIdentifier.parse(collectionName)); + } + + @Override + public boolean isFirstClassColumn(String fieldName) { + return columnMappings.containsKey(fieldName); + } + + @Override + public Optional getColumnType(String fieldName) { + PostgresColumnInfo columnInfo = columnMappings.get(fieldName); + return columnInfo != null ? Optional.of(columnInfo.getColumnType()) : Optional.empty(); + } + + @Override + public Optional getColumnName(String fieldName) { + PostgresColumnInfo columnInfo = columnMappings.get(fieldName); + return columnInfo != null ? Optional.of(columnInfo.getColumnName()) : Optional.empty(); + } + + @Override + public Set getAllFirstClassColumns() { + return Collections.unmodifiableSet(columnMappings.keySet()); + } + + @Override + public String getTableName() { + return tableName; + } + + /** + * Builds the column mappings by querying the database schema. + * + * @param connection the database connection + * @param tableIdentifier the table identifier + * @return a map of field names to PostgresColumnInfo objects + * @throws SQLException if there's an error querying the database + */ + private Map buildColumnMappings( + Connection connection, PostgresTableIdentifier tableIdentifier) throws SQLException { + + Map mappings = new HashMap<>(); + + String query = + "SELECT column_name, data_type, udt_name " + + "FROM information_schema.columns " + + "WHERE table_name = ? AND table_schema = ? " + + "ORDER BY ordinal_position"; + + String schemaName = tableIdentifier.getSchema().orElse("public"); + String tableName = tableIdentifier.getTableName(); + + try (PreparedStatement stmt = connection.prepareStatement(query)) { + stmt.setString(1, tableName); + stmt.setString(2, schemaName); + + log.debug("Querying schema for table: {}.{}", schemaName, tableName); + + try (ResultSet rs = stmt.executeQuery()) { + if (rs != null) { + while (rs.next()) { + String columnName = rs.getString("column_name"); + String dataType = rs.getString("data_type"); + String udtName = rs.getString("udt_name"); + + PostgresColumnType columnType = PostgresColumnType.fromPostgresType(dataType, udtName); + + // Only include first-class columns (non-JSONB) in the registry + if (columnType.isFirstClassField()) { + PostgresColumnInfo columnInfo = new PostgresColumnInfo(columnName, columnType); + mappings.put(columnName, columnInfo); + + log.debug( + "Registered first-class column: {} -> {} ({})", columnName, columnType, dataType); + } else { + log.debug("Skipping JSONB column: {} ({})", columnName, dataType); + } + } + } + } + } catch (SQLException e) { + log.error( + "Failed to query schema for table {}.{}: {}", schemaName, tableName, e.getMessage()); + throw e; + } + + return mappings; + } + + @Override + public String toString() { + return String.format( + "PostgresColumnRegistryImpl{tableName='%s', firstClassColumns=%d}", + tableName, columnMappings.size()); + } +} diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/registry/PostgresColumnType.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/registry/PostgresColumnType.java new file mode 100644 index 00000000..1ad53e09 --- /dev/null +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/registry/PostgresColumnType.java @@ -0,0 +1,88 @@ +package org.hypertrace.core.documentstore.postgres.registry; + +/** + * Enumeration of PostgreSQL column types supported by the document store. This enum maps PostgreSQL + * data types to their corresponding Java representations and determines the appropriate query + * generation strategy. + */ +public enum PostgresColumnType { + /** PostgreSQL text/varchar types - mapped to Java String */ + TEXT, + + /** PostgreSQL bigint/int8 types - mapped to Java Long */ + BIGINT, + + /** PostgreSQL double precision/float8 types - mapped to Java Double */ + DOUBLE_PRECISION, + + /** PostgreSQL boolean/bool types - mapped to Java Boolean */ + BOOLEAN, + + /** PostgreSQL text array types - mapped to Java String[] */ + TEXT_ARRAY, + + /** PostgreSQL jsonb/json types - existing document storage format */ + JSONB; + + /** + * Determines if this column type represents a first-class (non-JSON) field. First-class fields + * are stored as native PostgreSQL types rather than within JSONB documents. + * + * @return true if this is a first-class field type, false if it's a JSON document field + */ + public boolean isFirstClassField() { + return this != JSONB; + } + + /** + * Maps PostgreSQL data type names to PostgresColumnType enum values. + * + * @param dataType the PostgreSQL data type name (e.g., "text", "bigint", "boolean") + * @param udtName the user-defined type name for arrays (e.g., "_text" for text arrays) + * @return the corresponding PostgresColumnType, or JSONB as fallback + */ + public static PostgresColumnType fromPostgresType(String dataType, String udtName) { + if (dataType == null) { + return JSONB; + } + + String lowerDataType = dataType.toLowerCase(); + + // Handle array types first (indicated by udtName starting with underscore) + if (udtName != null && udtName.startsWith("_")) { + switch (udtName.toLowerCase()) { + case "_text": + return TEXT_ARRAY; + default: + return JSONB; // Unsupported array type, fallback to JSONB + } + } + + // Handle scalar types + switch (lowerDataType) { + case "text": + case "varchar": + case "character varying": + return TEXT; + + case "bigint": + case "int8": + return BIGINT; + + case "double precision": + case "float8": + return DOUBLE_PRECISION; + + case "boolean": + case "bool": + return BOOLEAN; + + case "jsonb": + case "json": + return JSONB; + + default: + return JSONB; // Unsupported type, fallback to existing JSONB behavior + } + } +}