From 00e9a9c5b421ef8a337abb1c5aeb30bfe5066307 Mon Sep 17 00:00:00 2001 From: Prashant Pandey Date: Sun, 28 Dec 2025 23:17:01 +0530 Subject: [PATCH 01/26] Added PostgresSchemaRegistry.java --- .../core/documentstore/Datastore.java | 6 + .../documentstore/commons/ColumnMetadata.java | 13 ++ .../documentstore/commons/SchemaRegistry.java | 12 ++ .../postgres/FlatPostgresCollection.java | 119 ++++++++++++++++- .../postgres/PostgresDatastore.java | 20 ++- .../postgres/PostgresMetadataFetcher.java | 126 ++++++++++++++++++ .../postgres/PostgresSchemaRegistry.java | 56 ++++++++ .../model/PostgresColumnMetadata.java | 39 ++++++ 8 files changed, 384 insertions(+), 7 deletions(-) create mode 100644 document-store/src/main/java/org/hypertrace/core/documentstore/commons/ColumnMetadata.java create mode 100644 document-store/src/main/java/org/hypertrace/core/documentstore/commons/SchemaRegistry.java create mode 100644 document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresMetadataFetcher.java create mode 100644 document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistry.java create mode 100644 document-store/src/main/java/org/hypertrace/core/documentstore/postgres/model/PostgresColumnMetadata.java diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/Datastore.java b/document-store/src/main/java/org/hypertrace/core/documentstore/Datastore.java index b426cb06..4022dd46 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/Datastore.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/Datastore.java @@ -2,6 +2,8 @@ import java.util.Map; import java.util.Set; +import org.hypertrace.core.documentstore.commons.ColumnMetadata; +import org.hypertrace.core.documentstore.commons.SchemaRegistry; import org.hypertrace.core.documentstore.metric.DocStoreMetricProvider; public interface Datastore { @@ -19,6 +21,10 @@ public interface Datastore { @SuppressWarnings("unused") DocStoreMetricProvider getDocStoreMetricProvider(); + default SchemaRegistry getSchemaRegistry() { + return null; + } + void close(); /** diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/commons/ColumnMetadata.java b/document-store/src/main/java/org/hypertrace/core/documentstore/commons/ColumnMetadata.java new file mode 100644 index 00000000..2850485a --- /dev/null +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/commons/ColumnMetadata.java @@ -0,0 +1,13 @@ +package org.hypertrace.core.documentstore.commons; + +import org.hypertrace.core.documentstore.expression.impl.DataType; + +public interface ColumnMetadata { + String getName(); + + DataType getCanonicalType(); + + String getInternalType(); + + boolean isNullable(); +} diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/commons/SchemaRegistry.java b/document-store/src/main/java/org/hypertrace/core/documentstore/commons/SchemaRegistry.java new file mode 100644 index 00000000..544743d5 --- /dev/null +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/commons/SchemaRegistry.java @@ -0,0 +1,12 @@ +package org.hypertrace.core.documentstore.commons; + +import java.util.Map; + +public interface SchemaRegistry { + + Map getSchema(String tableName); + + void invalidate(String tableName); + + T getColumnOrRefresh(String tableName, String colName); +} diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java index 0b7afc62..99b9fe0e 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java @@ -1,10 +1,18 @@ package org.hypertrace.core.documentstore.postgres; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; import java.io.IOException; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.util.ArrayList; +import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.Optional; import java.util.Set; +import java.util.stream.Collectors; import org.hypertrace.core.documentstore.BulkArrayValueUpdateRequest; import org.hypertrace.core.documentstore.BulkDeleteResult; import org.hypertrace.core.documentstore.BulkUpdateRequest; @@ -18,6 +26,7 @@ import org.hypertrace.core.documentstore.model.options.QueryOptions; import org.hypertrace.core.documentstore.model.options.UpdateOptions; import org.hypertrace.core.documentstore.model.subdoc.SubDocumentUpdate; +import org.hypertrace.core.documentstore.postgres.model.PostgresColumnMetadata; import org.hypertrace.core.documentstore.postgres.query.v1.PostgresQueryParser; import org.hypertrace.core.documentstore.postgres.query.v1.transformer.FlatPostgresFieldTransformer; import org.hypertrace.core.documentstore.query.Query; @@ -34,11 +43,18 @@ public class FlatPostgresCollection extends PostgresCollection { private static final Logger LOGGER = LoggerFactory.getLogger(FlatPostgresCollection.class); + private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); private static final String WRITE_NOT_SUPPORTED = "Write operations are not supported for flat collections yet!"; - FlatPostgresCollection(final PostgresClient client, final String collectionName) { + private final PostgresSchemaRegistry schemaRegistry; + + FlatPostgresCollection( + final PostgresClient client, + final String collectionName, + final PostgresSchemaRegistry schemaRegistry) { super(client, collectionName); + this.schemaRegistry = schemaRegistry; } @Override @@ -81,7 +97,106 @@ public boolean upsert(Key key, Document document) throws IOException { @Override public Document upsertAndReturn(Key key, Document document) throws IOException { - throw new UnsupportedOperationException(WRITE_NOT_SUPPORTED); + String tableName = tableIdentifier.getTableName(); + Map schema = schemaRegistry.getSchema(tableName); + + if (schema.isEmpty()) { + throw new IOException("No schema found for table: " + tableName); + } + + try { + JsonNode docJson = OBJECT_MAPPER.readTree(document.toJson()); + List columns = new ArrayList<>(); + List values = new ArrayList<>(); + + // Extract fields from document that exist in schema + Iterator> fields = docJson.fields(); + while (fields.hasNext()) { + Map.Entry field = fields.next(); + String colName = field.getKey(); + PostgresColumnMetadata colMeta = schemaRegistry.getColumnOrRefresh(tableName, colName); + if (colMeta != null) { + columns.add(colName); + values.add(extractValue(field.getValue())); + } + } + + if (columns.isEmpty()) { + throw new IOException("No matching columns found in schema for document"); + } + + // Build UPSERT SQL: INSERT ... ON CONFLICT DO UPDATE + String columnList = String.join(", ", columns); + String placeholders = columns.stream().map(c -> "?").collect(Collectors.joining(", ")); + String updateSet = columns.stream() + .map(c -> c + " = EXCLUDED." + c) + .collect(Collectors.joining(", ")); + + // Determine primary key column (assume first column or 'id') + String pkColumn = schema.containsKey("id") ? "id" : columns.get(0); + + String sql = String.format( + "INSERT INTO %s (%s) VALUES (%s) ON CONFLICT (%s) DO UPDATE SET %s RETURNING *", + tableIdentifier, columnList, placeholders, pkColumn, updateSet); + + try (PreparedStatement ps = client.getConnection().prepareStatement(sql)) { + for (int i = 0; i < values.size(); i++) { + ps.setObject(i + 1, values.get(i)); + } + + try (ResultSet rs = ps.executeQuery()) { + if (rs.next()) { + return resultSetToDocument(rs, columns); + } + } + } + return document; + } catch (SQLException e) { + LOGGER.error("SQLException in upsertAndReturn. key: {} document: {}", key, document, e); + throw new IOException(e); + } + } + + private Object extractValue(JsonNode node) { + if (node.isNull()) { + return null; + } else if (node.isBoolean()) { + return node.booleanValue(); + } else if (node.isInt()) { + return node.intValue(); + } else if (node.isLong()) { + return node.longValue(); + } else if (node.isDouble() || node.isFloat()) { + return node.doubleValue(); + } else if (node.isTextual()) { + return node.textValue(); + } else { + return node.toString(); + } + } + + private Document resultSetToDocument(ResultSet rs, List columns) + throws SQLException, IOException { + StringBuilder json = new StringBuilder("{"); + for (int i = 0; i < columns.size(); i++) { + if (i > 0) { + json.append(","); + } + String col = columns.get(i); + Object value = rs.getObject(col); + json.append("\"").append(col).append("\":"); + if (value == null) { + json.append("null"); + } else if (value instanceof String) { + json.append("\"").append(value).append("\""); + } else if (value instanceof Boolean) { + json.append(value); + } else { + json.append(value); + } + } + json.append("}"); + return new org.hypertrace.core.documentstore.JSONDocument(json.toString()); } @Override 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 afbb6f98..4b740322 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 @@ -20,11 +20,14 @@ import org.hypertrace.core.documentstore.Collection; import org.hypertrace.core.documentstore.Datastore; import org.hypertrace.core.documentstore.DocumentType; +import org.hypertrace.core.documentstore.commons.ColumnMetadata; +import org.hypertrace.core.documentstore.commons.SchemaRegistry; import org.hypertrace.core.documentstore.metric.DocStoreMetricProvider; import org.hypertrace.core.documentstore.metric.postgres.PostgresDocStoreMetricProvider; 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.model.PostgresColumnMetadata; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -37,6 +40,7 @@ public class PostgresDatastore implements Datastore { private final PostgresClient client; private final String database; private final DocStoreMetricProvider docStoreMetricProvider; + private final SchemaRegistry schemaRegistry; public PostgresDatastore(@NonNull final DatastoreConfig datastoreConfig) { final ConnectionConfig connectionConfig = datastoreConfig.connectionConfig(); @@ -57,6 +61,7 @@ public PostgresDatastore(@NonNull final DatastoreConfig datastoreConfig) { client = new PostgresClient(postgresConnectionConfig); database = connectionConfig.database(); docStoreMetricProvider = new PostgresDocStoreMetricProvider(this, postgresConnectionConfig); + schemaRegistry = new PostgresSchemaRegistry(new PostgresMetadataFetcher(this)); } catch (final IllegalArgumentException e) { throw new IllegalArgumentException( String.format("Unable to instantiate PostgresClient with config:%s", connectionConfig), @@ -81,7 +86,7 @@ public Set listCollections() { Set collections = new HashSet<>(); try { DatabaseMetaData metaData = client.getConnection().getMetaData(); - ResultSet tables = metaData.getTables(null, null, "%", new String[] {"TABLE"}); + ResultSet tables = metaData.getTables(null, null, "%", new String[]{"TABLE"}); while (tables.next()) { Optional nonPublicSchema = Optional.ofNullable(tables.getString("TABLE_SCHEM")) @@ -161,10 +166,9 @@ public Collection getCollection(String collectionName) { @Override public Collection getCollectionForType(String collectionName, DocumentType documentType) { switch (documentType) { - case FLAT: - { - return new FlatPostgresCollection(client, collectionName); - } + case FLAT: { + return new FlatPostgresCollection(client, collectionName, (PostgresSchemaRegistry) schemaRegistry); + } case NESTED: return getCollection(collectionName); default: @@ -189,6 +193,12 @@ public DocStoreMetricProvider getDocStoreMetricProvider() { return docStoreMetricProvider; } + @SuppressWarnings("unchecked") + @Override + public SchemaRegistry getSchemaRegistry() { + return schemaRegistry; + } + @Override public void close() { try { diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresMetadataFetcher.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresMetadataFetcher.java new file mode 100644 index 00000000..b2e2e957 --- /dev/null +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresMetadataFetcher.java @@ -0,0 +1,126 @@ +package org.hypertrace.core.documentstore.postgres; + +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.util.HashMap; +import java.util.Map; +import lombok.AllArgsConstructor; +import org.hypertrace.core.documentstore.expression.impl.DataType; +import org.hypertrace.core.documentstore.postgres.model.PostgresColumnMetadata; +import org.hypertrace.core.documentstore.postgres.query.v1.parser.filter.nonjson.field.PostgresDataType; + +/** + * Fetches schema metadata directly from Postgres system catalogs. Hardcoded to query + * information_schema.columns. + */ +@AllArgsConstructor +public class PostgresMetadataFetcher { + + private final PostgresDatastore datastore; + + // Hardcoded SQL for high-performance schema discovery + private static final String DISCOVERY_SQL = + "SELECT column_name, udt_name, is_nullable " + + "FROM information_schema.columns " + + "WHERE table_schema = 'public' AND table_name = ?"; + + public Map fetch(String tableName) { + Map metadataMap = new HashMap<>(); + + try (Connection conn = datastore.getPostgresClient(); + PreparedStatement ps = conn.prepareStatement(DISCOVERY_SQL)) { + + ps.setString(1, tableName); + try (ResultSet rs = ps.executeQuery()) { + while (rs.next()) { + String columnName = rs.getString("column_name"); + String udtName = rs.getString("udt_name"); + boolean isNullable = "YES".equalsIgnoreCase(rs.getString("is_nullable")); + metadataMap.put( + columnName, + new PostgresColumnMetadata( + columnName, + mapToCanonicalType(udtName), + mapToPostgresType(udtName), + udtName, + isNullable)); + } + } + return metadataMap; + } catch (SQLException e) { + throw new RuntimeException("Failed to fetch Postgres metadata for table: " + tableName, e); + } + } + + /** + * Maps Postgres udt_name to canonical DataType. + */ + private DataType mapToCanonicalType(String udtName) { + if (udtName == null) { + return DataType.UNSPECIFIED; + } + + switch (udtName.toLowerCase()) { + case "int4": + case "int2": + return DataType.INTEGER; + case "int8": + return DataType.LONG; + case "float4": + return DataType.FLOAT; + case "float8": + case "numeric": + return DataType.DOUBLE; + case "bool": + return DataType.BOOLEAN; + case "timestamptz": + return DataType.TIMESTAMPTZ; + case "date": + return DataType.DATE; + case "text": + case "varchar": + case "bpchar": + case "uuid": + return DataType.STRING; + default: + return DataType.UNSPECIFIED; + } + } + + /** + * Maps Postgres udt_name to PostgresDataType. + */ + private PostgresDataType mapToPostgresType(String udtName) { + if (udtName == null) { + return PostgresDataType.UNKNOWN; + } + + switch (udtName.toLowerCase()) { + case "int4": + case "int2": + return PostgresDataType.INTEGER; + case "int8": + return PostgresDataType.BIGINT; + case "float4": + return PostgresDataType.REAL; + case "float8": + case "numeric": + return PostgresDataType.DOUBLE_PRECISION; + case "bool": + return PostgresDataType.BOOLEAN; + case "timestamptz": + return PostgresDataType.TIMESTAMPTZ; + case "date": + return PostgresDataType.DATE; + case "text": + case "varchar": + case "bpchar": + case "uuid": + return PostgresDataType.TEXT; + default: + return PostgresDataType.UNKNOWN; + } + } +} diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistry.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistry.java new file mode 100644 index 00000000..bc5ad46e --- /dev/null +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistry.java @@ -0,0 +1,56 @@ +package org.hypertrace.core.documentstore.postgres; + +import com.google.common.cache.CacheBuilder; +import com.google.common.cache.CacheLoader; +import com.google.common.cache.LoadingCache; +import java.util.Map; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeUnit; +import org.hypertrace.core.documentstore.commons.SchemaRegistry; +import org.hypertrace.core.documentstore.postgres.model.PostgresColumnMetadata; + +public class PostgresSchemaRegistry implements SchemaRegistry { + + private final LoadingCache> cache; + private final PostgresMetadataFetcher fetcher; + + public PostgresSchemaRegistry(PostgresMetadataFetcher fetcher) { + this.fetcher = fetcher; + this.cache = + CacheBuilder.newBuilder() + .expireAfterWrite(24, TimeUnit.HOURS) + .build( + new CacheLoader<>() { + @Override + public Map load(String tableName) { + return fetcher.fetch(tableName); // Hardcoded SQL Discovery + } + }); + } + + @Override + public Map getSchema(String tableName) { + try { + return cache.get(tableName); + } catch (ExecutionException e) { + throw new RuntimeException("Failed to fetch schema for " + tableName, e.getCause()); + } + } + + @Override + public void invalidate(String tableName) { + cache.invalidate(tableName); + } + + @Override + public PostgresColumnMetadata getColumnOrRefresh(String tableName, String colName) { + Map schema = getSchema(tableName); + + if (!schema.containsKey(colName)) { + invalidate(tableName); + schema = getSchema(tableName); + } + + return schema.get(colName); + } +} diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/model/PostgresColumnMetadata.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/model/PostgresColumnMetadata.java new file mode 100644 index 00000000..7f2a18fd --- /dev/null +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/model/PostgresColumnMetadata.java @@ -0,0 +1,39 @@ +package org.hypertrace.core.documentstore.postgres.model; + +import lombok.AllArgsConstructor; +import lombok.Builder; +import lombok.Getter; +import org.hypertrace.core.documentstore.commons.ColumnMetadata; +import org.hypertrace.core.documentstore.expression.impl.DataType; +import org.hypertrace.core.documentstore.postgres.query.v1.parser.filter.nonjson.field.PostgresDataType; + +@Builder +@AllArgsConstructor +public class PostgresColumnMetadata implements ColumnMetadata { + + private final String colName; + private final DataType canonicalType; + @Getter private final PostgresDataType postgresType; + private final String pgType; + private final boolean nullable; + + @Override + public String getName() { + return colName; + } + + @Override + public DataType getCanonicalType() { + return canonicalType; + } + + @Override + public String getInternalType() { + return pgType; + } + + @Override + public boolean isNullable() { + return nullable; + } +} From 31846e946d90dd7e3e525d7fc769e152bc575a59 Mon Sep 17 00:00:00 2001 From: Prashant Pandey Date: Sun, 28 Dec 2025 23:17:12 +0530 Subject: [PATCH 02/26] Spotless --- .../postgres/FlatPostgresCollection.java | 12 ++++++------ .../documentstore/postgres/PostgresDatastore.java | 11 ++++++----- .../postgres/PostgresMetadataFetcher.java | 8 ++------ 3 files changed, 14 insertions(+), 17 deletions(-) diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java index 99b9fe0e..fa1c1ea1 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java @@ -128,16 +128,16 @@ public Document upsertAndReturn(Key key, Document document) throws IOException { // Build UPSERT SQL: INSERT ... ON CONFLICT DO UPDATE String columnList = String.join(", ", columns); String placeholders = columns.stream().map(c -> "?").collect(Collectors.joining(", ")); - String updateSet = columns.stream() - .map(c -> c + " = EXCLUDED." + c) - .collect(Collectors.joining(", ")); + String updateSet = + columns.stream().map(c -> c + " = EXCLUDED." + c).collect(Collectors.joining(", ")); // Determine primary key column (assume first column or 'id') String pkColumn = schema.containsKey("id") ? "id" : columns.get(0); - String sql = String.format( - "INSERT INTO %s (%s) VALUES (%s) ON CONFLICT (%s) DO UPDATE SET %s RETURNING *", - tableIdentifier, columnList, placeholders, pkColumn, updateSet); + String sql = + String.format( + "INSERT INTO %s (%s) VALUES (%s) ON CONFLICT (%s) DO UPDATE SET %s RETURNING *", + tableIdentifier, columnList, placeholders, pkColumn, updateSet); try (PreparedStatement ps = client.getConnection().prepareStatement(sql)) { for (int i = 0; i < values.size(); i++) { 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 4b740322..9a7fa206 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 @@ -20,7 +20,6 @@ import org.hypertrace.core.documentstore.Collection; import org.hypertrace.core.documentstore.Datastore; import org.hypertrace.core.documentstore.DocumentType; -import org.hypertrace.core.documentstore.commons.ColumnMetadata; import org.hypertrace.core.documentstore.commons.SchemaRegistry; import org.hypertrace.core.documentstore.metric.DocStoreMetricProvider; import org.hypertrace.core.documentstore.metric.postgres.PostgresDocStoreMetricProvider; @@ -86,7 +85,7 @@ public Set listCollections() { Set collections = new HashSet<>(); try { DatabaseMetaData metaData = client.getConnection().getMetaData(); - ResultSet tables = metaData.getTables(null, null, "%", new String[]{"TABLE"}); + ResultSet tables = metaData.getTables(null, null, "%", new String[] {"TABLE"}); while (tables.next()) { Optional nonPublicSchema = Optional.ofNullable(tables.getString("TABLE_SCHEM")) @@ -166,9 +165,11 @@ public Collection getCollection(String collectionName) { @Override public Collection getCollectionForType(String collectionName, DocumentType documentType) { switch (documentType) { - case FLAT: { - return new FlatPostgresCollection(client, collectionName, (PostgresSchemaRegistry) schemaRegistry); - } + case FLAT: + { + return new FlatPostgresCollection( + client, collectionName, (PostgresSchemaRegistry) schemaRegistry); + } case NESTED: return getCollection(collectionName); default: diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresMetadataFetcher.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresMetadataFetcher.java index b2e2e957..474df319 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresMetadataFetcher.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresMetadataFetcher.java @@ -54,9 +54,7 @@ public Map fetch(String tableName) { } } - /** - * Maps Postgres udt_name to canonical DataType. - */ + /** Maps Postgres udt_name to canonical DataType. */ private DataType mapToCanonicalType(String udtName) { if (udtName == null) { return DataType.UNSPECIFIED; @@ -89,9 +87,7 @@ private DataType mapToCanonicalType(String udtName) { } } - /** - * Maps Postgres udt_name to PostgresDataType. - */ + /** Maps Postgres udt_name to PostgresDataType. */ private PostgresDataType mapToPostgresType(String udtName) { if (udtName == null) { return PostgresDataType.UNKNOWN; From 2fdbf0eb7c3d34051c4e60a129058c300f8c981a Mon Sep 17 00:00:00 2001 From: Prashant Pandey Date: Mon, 29 Dec 2025 14:41:02 +0530 Subject: [PATCH 03/26] WIP --- .../postgres/PostgresSchemaRegistry.java | 142 +++++++- .../postgres/PostgresSchemaRegistryTest.java | 338 ++++++++++++++++++ 2 files changed, 475 insertions(+), 5 deletions(-) create mode 100644 document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistryTest.java diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistry.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistry.java index bc5ad46e..f3c06673 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistry.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistry.java @@ -3,31 +3,124 @@ import com.google.common.cache.CacheBuilder; import com.google.common.cache.CacheLoader; import com.google.common.cache.LoadingCache; +import java.time.Clock; +import java.time.Duration; +import java.time.Instant; import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; import org.hypertrace.core.documentstore.commons.SchemaRegistry; import org.hypertrace.core.documentstore.postgres.model.PostgresColumnMetadata; +/** + * A lazily-loaded, cached schema registry for PostgreSQL tables. + * + *

This registry fetches and caches column metadata from PostgreSQL's {@code information_schema} + * on demand. It provides: + * + *

    + *
  • Lazy loading: Schema metadata is fetched only when first requested for the + * particular table. + *
  • TTL-based caching: Cached schemas expire after a configurable duration (default: 24 + * hours). + *
  • Circuit breaker: Prevents excessive database calls by enforcing a cooldown period + * between refresh attempts for missing columns (default: 15 minutes). + *
+ * + *

Usage Example

+ * + * 
{@code
+ * PostgresMetadataFetcher fetcher = new PostgresMetadataFetcher(datastore);
+ * PostgresSchemaRegistry registry = new PostgresSchemaRegistry(fetcher);
+ *
+ * // Get all columns for a table
+ * Map schema = registry.getSchema("my_table");
+ *
+ * // Get a specific column, refreshing if not found (respects cooldown)
+ * PostgresColumnMetadata column = registry.getColumnOrRefresh("my_table", "my_column");
+ * }
+ * + *

Thread Safety

+ * + *

This class is thread-safe. The underlying Guava {@link LoadingCache} and {@link + * ConcurrentHashMap} handle concurrent access. + * + * @see PostgresMetadataFetcher + * @see PostgresColumnMetadata + */ public class PostgresSchemaRegistry implements SchemaRegistry { + /** Default cache expiry time: 24 hours. */ + private static final Duration DEFAULT_CACHE_EXPIRY = Duration.ofHours(24); + + /** Default cooldown period between refresh attempts: 15 minutes. */ + private static final Duration DEFAULT_REFRESH_COOLDOWN = Duration.ofMinutes(15); + private final LoadingCache> cache; - private final PostgresMetadataFetcher fetcher; + private final Map lastRefreshTimes; + private final Duration refreshCooldown; + private final Clock clock; + /** + * Creates a new schema registry with default settings. + * + *

Uses default cache expiry of 24 hours and refresh cooldown of 15 minutes. + * + * @param fetcher the metadata fetcher to use for loading schema information + */ public PostgresSchemaRegistry(PostgresMetadataFetcher fetcher) { - this.fetcher = fetcher; + this(fetcher, DEFAULT_CACHE_EXPIRY, DEFAULT_REFRESH_COOLDOWN, Clock.systemUTC()); + } + + /** + * Creates a new schema registry with custom cache settings. + * + * @param fetcher the metadata fetcher to use for loading schema information + * @param cacheExpiry how long to keep cached schemas before they expire + * @param refreshCooldown minimum time between refresh attempts for missing columns + */ + public PostgresSchemaRegistry( + PostgresMetadataFetcher fetcher, Duration cacheExpiry, Duration refreshCooldown) { + this(fetcher, cacheExpiry, refreshCooldown, Clock.systemUTC()); + } + + /** + * Creates a new schema registry with custom settings and clock (for testing). + * + * @param fetcher the metadata fetcher to use for loading schema information + * @param cacheExpiry how long to keep cached schemas before they expire + * @param refreshCooldown minimum time between refresh attempts for missing columns + * @param clock the clock to use for time-based operations + */ + PostgresSchemaRegistry( + PostgresMetadataFetcher fetcher, + Duration cacheExpiry, + Duration refreshCooldown, + Clock clock) { + this.refreshCooldown = refreshCooldown; + this.clock = clock; + this.lastRefreshTimes = new ConcurrentHashMap<>(); this.cache = CacheBuilder.newBuilder() - .expireAfterWrite(24, TimeUnit.HOURS) + .expireAfterWrite(cacheExpiry.toMinutes(), TimeUnit.MINUTES) .build( new CacheLoader<>() { @Override public Map load(String tableName) { - return fetcher.fetch(tableName); // Hardcoded SQL Discovery + lastRefreshTimes.put(tableName, clock.instant()); + return fetcher.fetch(tableName); } }); } + /** + * Gets the schema for a table, loading it from the database if not cached. + * + * @param tableName the name of the table + * @return a map of column names to their metadata + * @throws RuntimeException if the schema cannot be fetched + */ @Override public Map getSchema(String tableName) { try { @@ -37,20 +130,59 @@ public Map getSchema(String tableName) { } } + /** + * Invalidates the cached schema for a specific table. + * + *

The next call to {@link #getSchema(String)} will reload the schema from the database. + * + * @param tableName the name of the table to invalidate + */ @Override public void invalidate(String tableName) { cache.invalidate(tableName); } + /** + * Gets metadata for a specific column, optionally refreshing the schema if the column is not + * found. + * + *

This method implements a circuit breaker pattern: + * + *

    + *
  • If the column exists in the cached schema, it is returned immediately. + *
  • If the column is not found and the cooldown period has elapsed since the last refresh, + * the schema is reloaded from the database. + *
  • If the column is not found but the cooldown period has not elapsed, {@code null} is + * returned without hitting the database. + *
+ * + * @param tableName the name of the table + * @param colName the name of the column + * @return the column metadata, or {@code null} if the column does not exist + */ @Override public PostgresColumnMetadata getColumnOrRefresh(String tableName, String colName) { Map schema = getSchema(tableName); - if (!schema.containsKey(colName)) { + if (!schema.containsKey(colName) && canRefresh(tableName)) { invalidate(tableName); schema = getSchema(tableName); } return schema.get(colName); } + + /** + * Checks if a refresh is allowed for the given table based on the cooldown period. + * + * @param tableName the name of the table + * @return {@code true} if the cooldown period has elapsed since the last refresh + */ + private boolean canRefresh(String tableName) { + Instant lastRefresh = lastRefreshTimes.get(tableName); + if (lastRefresh == null) { + return true; + } + return Duration.between(lastRefresh, clock.instant()).compareTo(refreshCooldown) >= 0; + } } diff --git a/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistryTest.java b/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistryTest.java new file mode 100644 index 00000000..6eb85b01 --- /dev/null +++ b/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistryTest.java @@ -0,0 +1,338 @@ +package org.hypertrace.core.documentstore.postgres; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.time.Clock; +import java.time.Duration; +import java.time.Instant; +import java.time.ZoneId; +import java.util.HashMap; +import java.util.Map; +import org.hypertrace.core.documentstore.expression.impl.DataType; +import org.hypertrace.core.documentstore.postgres.model.PostgresColumnMetadata; +import org.hypertrace.core.documentstore.postgres.query.v1.parser.filter.nonjson.field.PostgresDataType; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; + +@ExtendWith(MockitoExtension.class) +class PostgresSchemaRegistryTest { + + private static final String TEST_TABLE = "test_table"; + private static final String COL_ID = "id"; + private static final String COL_NAME = "name"; + private static final String COL_PRICE = "price"; + private static final Duration CACHE_EXPIRY = Duration.ofHours(24); + private static final Duration REFRESH_COOLDOWN = Duration.ofMinutes(15); + + @Mock private PostgresMetadataFetcher fetcher; + + private PostgresSchemaRegistry registry; + private MutableClock mutableClock; + + @BeforeEach + void setUp() { + mutableClock = new MutableClock(Instant.parse("2024-01-01T00:00:00Z")); + registry = new PostgresSchemaRegistry(fetcher, CACHE_EXPIRY, REFRESH_COOLDOWN, mutableClock); + } + + @Test + void getSchemaLoadsFromFetcherOnCacheMiss() { + Map expectedSchema = createTestSchema(); + when(fetcher.fetch(TEST_TABLE)).thenReturn(expectedSchema); + + Map result = registry.getSchema(TEST_TABLE); + + assertEquals(expectedSchema, result); + verify(fetcher, times(1)).fetch(TEST_TABLE); + } + + @Test + void getSchemaReturnsCachedValueOnSubsequentCalls() { + Map expectedSchema = createTestSchema(); + when(fetcher.fetch(TEST_TABLE)).thenReturn(expectedSchema); + + // First call - loads from fetcher + Map result1 = registry.getSchema(TEST_TABLE); + // Second call - should use cache + Map result2 = registry.getSchema(TEST_TABLE); + + assertEquals(expectedSchema, result1); + assertEquals(expectedSchema, result2); + // Fetcher should only be called once + verify(fetcher, times(1)).fetch(TEST_TABLE); + } + + @Test + void getSchemaLoadsEachTableIndependently() { + String table1 = "table1"; + String table2 = "table2"; + Map schema1 = createTestSchema(); + Map schema2 = new HashMap<>(); + schema2.put( + "other_col", + PostgresColumnMetadata.builder() + .colName("other_col") + .canonicalType(DataType.BOOLEAN) + .postgresType(PostgresDataType.BOOLEAN) + .pgType("bool") + .nullable(false) + .build()); + + when(fetcher.fetch(table1)).thenReturn(schema1); + when(fetcher.fetch(table2)).thenReturn(schema2); + + Map result1 = registry.getSchema(table1); + Map result2 = registry.getSchema(table2); + + assertEquals(schema1, result1); + assertEquals(schema2, result2); + verify(fetcher, times(1)).fetch(table1); + verify(fetcher, times(1)).fetch(table2); + } + + @Test + void invalidateClearsSpecificTableCache() { + Map schema = createTestSchema(); + when(fetcher.fetch(TEST_TABLE)).thenReturn(schema); + + // Load into cache + registry.getSchema(TEST_TABLE); + verify(fetcher, times(1)).fetch(TEST_TABLE); + + // Invalidate the cache + registry.invalidate(TEST_TABLE); + + // Next call should reload from fetcher + registry.getSchema(TEST_TABLE); + verify(fetcher, times(2)).fetch(TEST_TABLE); + } + + @Test + void invalidateDoesNotAffectOtherTables() { + String table1 = "table1"; + String table2 = "table2"; + Map schema1 = createTestSchema(); + Map schema2 = new HashMap<>(); + + when(fetcher.fetch(table1)).thenReturn(schema1); + when(fetcher.fetch(table2)).thenReturn(schema2); + + // Load both tables into cache + registry.getSchema(table1); + registry.getSchema(table2); + + // Invalidate only table1 + registry.invalidate(table1); + + // table1 should reload, table2 should use cache + registry.getSchema(table1); + registry.getSchema(table2); + + verify(fetcher, times(2)).fetch(table1); + verify(fetcher, times(1)).fetch(table2); + } + + @Test + void getColumnOrRefreshReturnsColumnIfExists() { + Map schema = createTestSchema(); + when(fetcher.fetch(TEST_TABLE)).thenReturn(schema); + + PostgresColumnMetadata result = registry.getColumnOrRefresh(TEST_TABLE, COL_ID); + + assertNotNull(result); + assertEquals(COL_ID, result.getName()); + assertEquals(DataType.INTEGER, result.getCanonicalType()); + // Should only call fetcher once (initial load) + verify(fetcher, times(1)).fetch(TEST_TABLE); + } + + @Test + void getColumnOrRefreshRefreshesSchemaIfColumnMissingAndCooldownExpired() { + // Initial schema without the "new_col" + Map initialSchema = createTestSchema(); + + // Updated schema with the "new_col" + Map updatedSchema = new HashMap<>(initialSchema); + updatedSchema.put( + "new_col", + PostgresColumnMetadata.builder() + .colName("new_col") + .canonicalType(DataType.STRING) + .postgresType(PostgresDataType.TEXT) + .pgType("text") + .nullable(true) + .build()); + + when(fetcher.fetch(TEST_TABLE)).thenReturn(initialSchema).thenReturn(updatedSchema); + + // First call loads the schema + registry.getSchema(TEST_TABLE); + verify(fetcher, times(1)).fetch(TEST_TABLE); + + // Advance time past cooldown period + mutableClock.advance(REFRESH_COOLDOWN.plusMinutes(1)); + + // Now try to get missing column - should trigger refresh + PostgresColumnMetadata result = registry.getColumnOrRefresh(TEST_TABLE, "new_col"); + + assertNotNull(result); + assertEquals("new_col", result.getName()); + // Should call fetcher twice (initial load + refresh after cooldown) + verify(fetcher, times(2)).fetch(TEST_TABLE); + } + + @Test + void getColumnOrRefreshDoesNotRefreshIfWithinCooldownPeriod() { + Map schema = createTestSchema(); + when(fetcher.fetch(TEST_TABLE)).thenReturn(schema); + + // First call loads the schema + registry.getSchema(TEST_TABLE); + verify(fetcher, times(1)).fetch(TEST_TABLE); + + // Try to get a missing column - should NOT refresh because we're within cooldown + PostgresColumnMetadata result = registry.getColumnOrRefresh(TEST_TABLE, "nonexistent_col"); + + assertNull(result); + // Should only call fetcher once (initial load, no refresh due to cooldown) + verify(fetcher, times(1)).fetch(TEST_TABLE); + } + + @Test + void getColumnOrRefreshRefreshesAfterCooldownExpires() { + Map schema = createTestSchema(); + when(fetcher.fetch(TEST_TABLE)).thenReturn(schema); + + // First call loads the schema + registry.getSchema(TEST_TABLE); + verify(fetcher, times(1)).fetch(TEST_TABLE); + + // Try immediately - should NOT refresh (within cooldown) + registry.getColumnOrRefresh(TEST_TABLE, "nonexistent_col"); + verify(fetcher, times(1)).fetch(TEST_TABLE); + + // Advance time past cooldown + mutableClock.advance(REFRESH_COOLDOWN.plusSeconds(1)); + + // Try again - should refresh now + registry.getColumnOrRefresh(TEST_TABLE, "nonexistent_col"); + verify(fetcher, times(2)).fetch(TEST_TABLE); + } + + @Test + void getColumnOrRefreshReturnsNullIfColumnStillMissingAfterRefresh() { + Map schema = createTestSchema(); + when(fetcher.fetch(TEST_TABLE)).thenReturn(schema); + + // First call loads the schema + registry.getSchema(TEST_TABLE); + + // Advance past cooldown + mutableClock.advance(REFRESH_COOLDOWN.plusMinutes(1)); + + // Try to get a column that doesn't exist even after refresh + PostgresColumnMetadata result = registry.getColumnOrRefresh(TEST_TABLE, "nonexistent_col"); + + assertNull(result); + // Should call fetcher twice (initial load + refresh attempt after cooldown) + verify(fetcher, times(2)).fetch(TEST_TABLE); + } + + @Test + void getColumnOrRefreshUsesExistingCacheBeforeRefresh() { + Map schema = createTestSchema(); + when(fetcher.fetch(TEST_TABLE)).thenReturn(schema); + + // Pre-populate cache + registry.getSchema(TEST_TABLE); + + // Get existing column - should not trigger refresh + PostgresColumnMetadata result = registry.getColumnOrRefresh(TEST_TABLE, COL_NAME); + + assertNotNull(result); + assertEquals(COL_NAME, result.getName()); + // Should only call fetcher once (initial getSchema) + verify(fetcher, times(1)).fetch(TEST_TABLE); + } + + @Test + void getSchemaThrowsRuntimeExceptionWhenFetcherFails() { + when(fetcher.fetch(TEST_TABLE)).thenThrow(new RuntimeException("Database error")); + + RuntimeException exception = + assertThrows(RuntimeException.class, () -> registry.getSchema(TEST_TABLE)); + + assertEquals("Database error", exception.getCause().getMessage()); + } + + private Map createTestSchema() { + Map schema = new HashMap<>(); + schema.put( + COL_ID, + PostgresColumnMetadata.builder() + .colName(COL_ID) + .canonicalType(DataType.INTEGER) + .postgresType(PostgresDataType.INTEGER) + .pgType("int4") + .nullable(false) + .build()); + schema.put( + COL_NAME, + PostgresColumnMetadata.builder() + .colName(COL_NAME) + .canonicalType(DataType.STRING) + .postgresType(PostgresDataType.TEXT) + .pgType("text") + .nullable(true) + .build()); + schema.put( + COL_PRICE, + PostgresColumnMetadata.builder() + .colName(COL_PRICE) + .canonicalType(DataType.DOUBLE) + .postgresType(PostgresDataType.DOUBLE_PRECISION) + .pgType("float8") + .nullable(true) + .build()); + return schema; + } + + /** A mutable clock for testing time-dependent behavior. */ + private static class MutableClock extends Clock { + private Instant currentInstant; + private final ZoneId zone; + + MutableClock(Instant initialInstant) { + this.currentInstant = initialInstant; + this.zone = ZoneId.of("UTC"); + } + + void advance(Duration duration) { + currentInstant = currentInstant.plus(duration); + } + + @Override + public ZoneId getZone() { + return zone; + } + + @Override + public Clock withZone(ZoneId zone) { + return this; + } + + @Override + public Instant instant() { + return currentInstant; + } + } +} From 1727dd006441af7cbec718460b3c02e72ee776e9 Mon Sep 17 00:00:00 2001 From: Prashant Pandey Date: Mon, 29 Dec 2025 15:04:53 +0530 Subject: [PATCH 04/26] Spotless --- .../documentstore/commons/SchemaRegistry.java | 3 +- .../postgres/FlatPostgresCollection.java | 3 +- .../postgres/PostgresMetadataFetcher.java | 1 - .../postgres/PostgresSchemaRegistry.java | 7 ++-- .../postgres/PostgresSchemaRegistryTest.java | 35 ++++++++++--------- 5 files changed, 26 insertions(+), 23 deletions(-) diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/commons/SchemaRegistry.java b/document-store/src/main/java/org/hypertrace/core/documentstore/commons/SchemaRegistry.java index 544743d5..68782184 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/commons/SchemaRegistry.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/commons/SchemaRegistry.java @@ -1,6 +1,7 @@ package org.hypertrace.core.documentstore.commons; import java.util.Map; +import java.util.Optional; public interface SchemaRegistry { @@ -8,5 +9,5 @@ public interface SchemaRegistry { void invalidate(String tableName); - T getColumnOrRefresh(String tableName, String colName); + Optional getColumnOrRefresh(String tableName, String colName); } diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java index fa1c1ea1..7798e016 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java @@ -114,8 +114,7 @@ public Document upsertAndReturn(Key key, Document document) throws IOException { while (fields.hasNext()) { Map.Entry field = fields.next(); String colName = field.getKey(); - PostgresColumnMetadata colMeta = schemaRegistry.getColumnOrRefresh(tableName, colName); - if (colMeta != null) { + if (schemaRegistry.getColumnOrRefresh(tableName, colName).isPresent()) { columns.add(colName); values.add(extractValue(field.getValue())); } diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresMetadataFetcher.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresMetadataFetcher.java index 474df319..c32842bc 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresMetadataFetcher.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresMetadataFetcher.java @@ -20,7 +20,6 @@ public class PostgresMetadataFetcher { private final PostgresDatastore datastore; - // Hardcoded SQL for high-performance schema discovery private static final String DISCOVERY_SQL = "SELECT column_name, udt_name, is_nullable " + "FROM information_schema.columns " diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistry.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistry.java index f3c06673..ff7d6164 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistry.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistry.java @@ -7,6 +7,7 @@ import java.time.Duration; import java.time.Instant; import java.util.Map; +import java.util.Optional; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; @@ -158,10 +159,10 @@ public void invalidate(String tableName) { * * @param tableName the name of the table * @param colName the name of the column - * @return the column metadata, or {@code null} if the column does not exist + * @return an Optional containing the column metadata, or empty if the column does not exist */ @Override - public PostgresColumnMetadata getColumnOrRefresh(String tableName, String colName) { + public Optional getColumnOrRefresh(String tableName, String colName) { Map schema = getSchema(tableName); if (!schema.containsKey(colName) && canRefresh(tableName)) { @@ -169,7 +170,7 @@ public PostgresColumnMetadata getColumnOrRefresh(String tableName, String colNam schema = getSchema(tableName); } - return schema.get(colName); + return Optional.ofNullable(schema.get(colName)); } /** diff --git a/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistryTest.java b/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistryTest.java index 6eb85b01..fa8e4a03 100644 --- a/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistryTest.java +++ b/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistryTest.java @@ -1,9 +1,9 @@ package org.hypertrace.core.documentstore.postgres; import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertNotNull; -import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -14,6 +14,7 @@ import java.time.ZoneId; import java.util.HashMap; import java.util.Map; +import java.util.Optional; import org.hypertrace.core.documentstore.expression.impl.DataType; import org.hypertrace.core.documentstore.postgres.model.PostgresColumnMetadata; import org.hypertrace.core.documentstore.postgres.query.v1.parser.filter.nonjson.field.PostgresDataType; @@ -146,11 +147,11 @@ void getColumnOrRefreshReturnsColumnIfExists() { Map schema = createTestSchema(); when(fetcher.fetch(TEST_TABLE)).thenReturn(schema); - PostgresColumnMetadata result = registry.getColumnOrRefresh(TEST_TABLE, COL_ID); + Optional result = registry.getColumnOrRefresh(TEST_TABLE, COL_ID); - assertNotNull(result); - assertEquals(COL_ID, result.getName()); - assertEquals(DataType.INTEGER, result.getCanonicalType()); + assertTrue(result.isPresent()); + assertEquals(COL_ID, result.get().getName()); + assertEquals(DataType.INTEGER, result.get().getCanonicalType()); // Should only call fetcher once (initial load) verify(fetcher, times(1)).fetch(TEST_TABLE); } @@ -182,10 +183,10 @@ void getColumnOrRefreshRefreshesSchemaIfColumnMissingAndCooldownExpired() { mutableClock.advance(REFRESH_COOLDOWN.plusMinutes(1)); // Now try to get missing column - should trigger refresh - PostgresColumnMetadata result = registry.getColumnOrRefresh(TEST_TABLE, "new_col"); + Optional result = registry.getColumnOrRefresh(TEST_TABLE, "new_col"); - assertNotNull(result); - assertEquals("new_col", result.getName()); + assertTrue(result.isPresent()); + assertEquals("new_col", result.get().getName()); // Should call fetcher twice (initial load + refresh after cooldown) verify(fetcher, times(2)).fetch(TEST_TABLE); } @@ -200,9 +201,10 @@ void getColumnOrRefreshDoesNotRefreshIfWithinCooldownPeriod() { verify(fetcher, times(1)).fetch(TEST_TABLE); // Try to get a missing column - should NOT refresh because we're within cooldown - PostgresColumnMetadata result = registry.getColumnOrRefresh(TEST_TABLE, "nonexistent_col"); + Optional result = + registry.getColumnOrRefresh(TEST_TABLE, "nonexistent_col"); - assertNull(result); + assertFalse(result.isPresent()); // Should only call fetcher once (initial load, no refresh due to cooldown) verify(fetcher, times(1)).fetch(TEST_TABLE); } @@ -240,9 +242,10 @@ void getColumnOrRefreshReturnsNullIfColumnStillMissingAfterRefresh() { mutableClock.advance(REFRESH_COOLDOWN.plusMinutes(1)); // Try to get a column that doesn't exist even after refresh - PostgresColumnMetadata result = registry.getColumnOrRefresh(TEST_TABLE, "nonexistent_col"); + Optional result = + registry.getColumnOrRefresh(TEST_TABLE, "nonexistent_col"); - assertNull(result); + assertFalse(result.isPresent()); // Should call fetcher twice (initial load + refresh attempt after cooldown) verify(fetcher, times(2)).fetch(TEST_TABLE); } @@ -256,10 +259,10 @@ void getColumnOrRefreshUsesExistingCacheBeforeRefresh() { registry.getSchema(TEST_TABLE); // Get existing column - should not trigger refresh - PostgresColumnMetadata result = registry.getColumnOrRefresh(TEST_TABLE, COL_NAME); + Optional result = registry.getColumnOrRefresh(TEST_TABLE, COL_NAME); - assertNotNull(result); - assertEquals(COL_NAME, result.getName()); + assertTrue(result.isPresent()); + assertEquals(COL_NAME, result.get().getName()); // Should only call fetcher once (initial getSchema) verify(fetcher, times(1)).fetch(TEST_TABLE); } From a62fbc218a6b3a68fd8dbd2d49e6a4b07d5ed39c Mon Sep 17 00:00:00 2001 From: Prashant Pandey Date: Mon, 29 Dec 2025 15:06:01 +0530 Subject: [PATCH 05/26] Remove unused method in SchemaRegistry --- .../hypertrace/core/documentstore/commons/SchemaRegistry.java | 2 -- 1 file changed, 2 deletions(-) diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/commons/SchemaRegistry.java b/document-store/src/main/java/org/hypertrace/core/documentstore/commons/SchemaRegistry.java index 68782184..ec3ea43a 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/commons/SchemaRegistry.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/commons/SchemaRegistry.java @@ -7,7 +7,5 @@ public interface SchemaRegistry { Map getSchema(String tableName); - void invalidate(String tableName); - Optional getColumnOrRefresh(String tableName, String colName); } From 6b7595bafcb036e553d8ee1780108cfeb42150d0 Mon Sep 17 00:00:00 2001 From: Prashant Pandey Date: Mon, 29 Dec 2025 15:07:11 +0530 Subject: [PATCH 06/26] Remove unused method in ColumnMetadata --- .../hypertrace/core/documentstore/commons/ColumnMetadata.java | 2 -- .../hypertrace/core/documentstore/commons/SchemaRegistry.java | 2 ++ 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/commons/ColumnMetadata.java b/document-store/src/main/java/org/hypertrace/core/documentstore/commons/ColumnMetadata.java index 2850485a..6fa89dc3 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/commons/ColumnMetadata.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/commons/ColumnMetadata.java @@ -7,7 +7,5 @@ public interface ColumnMetadata { DataType getCanonicalType(); - String getInternalType(); - boolean isNullable(); } diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/commons/SchemaRegistry.java b/document-store/src/main/java/org/hypertrace/core/documentstore/commons/SchemaRegistry.java index ec3ea43a..68782184 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/commons/SchemaRegistry.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/commons/SchemaRegistry.java @@ -7,5 +7,7 @@ public interface SchemaRegistry { Map getSchema(String tableName); + void invalidate(String tableName); + Optional getColumnOrRefresh(String tableName, String colName); } From 7b4ef2a7e2eeef4d8d184c935dc4182d0ecd883d Mon Sep 17 00:00:00 2001 From: Prashant Pandey Date: Mon, 29 Dec 2025 15:07:26 +0530 Subject: [PATCH 07/26] WIP --- .../documentstore/postgres/model/PostgresColumnMetadata.java | 5 ----- 1 file changed, 5 deletions(-) diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/model/PostgresColumnMetadata.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/model/PostgresColumnMetadata.java index 7f2a18fd..508b3dc6 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/model/PostgresColumnMetadata.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/model/PostgresColumnMetadata.java @@ -27,11 +27,6 @@ public DataType getCanonicalType() { return canonicalType; } - @Override - public String getInternalType() { - return pgType; - } - @Override public boolean isNullable() { return nullable; From 598cb25e9e65a02e2e18aa332834a711ce622ff1 Mon Sep 17 00:00:00 2001 From: Prashant Pandey Date: Mon, 29 Dec 2025 15:28:47 +0530 Subject: [PATCH 08/26] WIP --- .../postgres/FlatPostgresCollection.java | 110 +----------------- .../postgres/PostgresDatastore.java | 2 +- .../postgres/PostgresMetadataFetcher.java | 4 +- 3 files changed, 4 insertions(+), 112 deletions(-) diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java index 7798e016..6dd89e1b 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java @@ -1,18 +1,10 @@ package org.hypertrace.core.documentstore.postgres; -import com.fasterxml.jackson.databind.JsonNode; -import com.fasterxml.jackson.databind.ObjectMapper; import java.io.IOException; -import java.sql.PreparedStatement; -import java.sql.ResultSet; -import java.sql.SQLException; -import java.util.ArrayList; -import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.Optional; import java.util.Set; -import java.util.stream.Collectors; import org.hypertrace.core.documentstore.BulkArrayValueUpdateRequest; import org.hypertrace.core.documentstore.BulkDeleteResult; import org.hypertrace.core.documentstore.BulkUpdateRequest; @@ -26,7 +18,6 @@ import org.hypertrace.core.documentstore.model.options.QueryOptions; import org.hypertrace.core.documentstore.model.options.UpdateOptions; import org.hypertrace.core.documentstore.model.subdoc.SubDocumentUpdate; -import org.hypertrace.core.documentstore.postgres.model.PostgresColumnMetadata; import org.hypertrace.core.documentstore.postgres.query.v1.PostgresQueryParser; import org.hypertrace.core.documentstore.postgres.query.v1.transformer.FlatPostgresFieldTransformer; import org.hypertrace.core.documentstore.query.Query; @@ -43,7 +34,6 @@ public class FlatPostgresCollection extends PostgresCollection { private static final Logger LOGGER = LoggerFactory.getLogger(FlatPostgresCollection.class); - private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); private static final String WRITE_NOT_SUPPORTED = "Write operations are not supported for flat collections yet!"; @@ -97,105 +87,7 @@ public boolean upsert(Key key, Document document) throws IOException { @Override public Document upsertAndReturn(Key key, Document document) throws IOException { - String tableName = tableIdentifier.getTableName(); - Map schema = schemaRegistry.getSchema(tableName); - - if (schema.isEmpty()) { - throw new IOException("No schema found for table: " + tableName); - } - - try { - JsonNode docJson = OBJECT_MAPPER.readTree(document.toJson()); - List columns = new ArrayList<>(); - List values = new ArrayList<>(); - - // Extract fields from document that exist in schema - Iterator> fields = docJson.fields(); - while (fields.hasNext()) { - Map.Entry field = fields.next(); - String colName = field.getKey(); - if (schemaRegistry.getColumnOrRefresh(tableName, colName).isPresent()) { - columns.add(colName); - values.add(extractValue(field.getValue())); - } - } - - if (columns.isEmpty()) { - throw new IOException("No matching columns found in schema for document"); - } - - // Build UPSERT SQL: INSERT ... ON CONFLICT DO UPDATE - String columnList = String.join(", ", columns); - String placeholders = columns.stream().map(c -> "?").collect(Collectors.joining(", ")); - String updateSet = - columns.stream().map(c -> c + " = EXCLUDED." + c).collect(Collectors.joining(", ")); - - // Determine primary key column (assume first column or 'id') - String pkColumn = schema.containsKey("id") ? "id" : columns.get(0); - - String sql = - String.format( - "INSERT INTO %s (%s) VALUES (%s) ON CONFLICT (%s) DO UPDATE SET %s RETURNING *", - tableIdentifier, columnList, placeholders, pkColumn, updateSet); - - try (PreparedStatement ps = client.getConnection().prepareStatement(sql)) { - for (int i = 0; i < values.size(); i++) { - ps.setObject(i + 1, values.get(i)); - } - - try (ResultSet rs = ps.executeQuery()) { - if (rs.next()) { - return resultSetToDocument(rs, columns); - } - } - } - return document; - } catch (SQLException e) { - LOGGER.error("SQLException in upsertAndReturn. key: {} document: {}", key, document, e); - throw new IOException(e); - } - } - - private Object extractValue(JsonNode node) { - if (node.isNull()) { - return null; - } else if (node.isBoolean()) { - return node.booleanValue(); - } else if (node.isInt()) { - return node.intValue(); - } else if (node.isLong()) { - return node.longValue(); - } else if (node.isDouble() || node.isFloat()) { - return node.doubleValue(); - } else if (node.isTextual()) { - return node.textValue(); - } else { - return node.toString(); - } - } - - private Document resultSetToDocument(ResultSet rs, List columns) - throws SQLException, IOException { - StringBuilder json = new StringBuilder("{"); - for (int i = 0; i < columns.size(); i++) { - if (i > 0) { - json.append(","); - } - String col = columns.get(i); - Object value = rs.getObject(col); - json.append("\"").append(col).append("\":"); - if (value == null) { - json.append("null"); - } else if (value instanceof String) { - json.append("\"").append(value).append("\""); - } else if (value instanceof Boolean) { - json.append(value); - } else { - json.append(value); - } - } - json.append("}"); - return new org.hypertrace.core.documentstore.JSONDocument(json.toString()); + throw new UnsupportedOperationException(WRITE_NOT_SUPPORTED); } @Override 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 9a7fa206..3b1fb345 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 @@ -60,7 +60,7 @@ public PostgresDatastore(@NonNull final DatastoreConfig datastoreConfig) { client = new PostgresClient(postgresConnectionConfig); database = connectionConfig.database(); docStoreMetricProvider = new PostgresDocStoreMetricProvider(this, postgresConnectionConfig); - schemaRegistry = new PostgresSchemaRegistry(new PostgresMetadataFetcher(this)); + schemaRegistry = new PostgresSchemaRegistry(new PostgresMetadataFetcher(client)); } catch (final IllegalArgumentException e) { throw new IllegalArgumentException( String.format("Unable to instantiate PostgresClient with config:%s", connectionConfig), diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresMetadataFetcher.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresMetadataFetcher.java index c32842bc..e61b26b7 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresMetadataFetcher.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresMetadataFetcher.java @@ -18,7 +18,7 @@ @AllArgsConstructor public class PostgresMetadataFetcher { - private final PostgresDatastore datastore; + private final PostgresClient client; private static final String DISCOVERY_SQL = "SELECT column_name, udt_name, is_nullable " @@ -28,7 +28,7 @@ public class PostgresMetadataFetcher { public Map fetch(String tableName) { Map metadataMap = new HashMap<>(); - try (Connection conn = datastore.getPostgresClient(); + try (Connection conn = client.getPooledConnection(); PreparedStatement ps = conn.prepareStatement(DISCOVERY_SQL)) { ps.setString(1, tableName); From 9c173b96e5a25c898a9b226070970df4867a5bba Mon Sep 17 00:00:00 2001 From: Prashant Pandey Date: Mon, 29 Dec 2025 15:45:23 +0530 Subject: [PATCH 09/26] Configure cache expiry and cooldown --- .../TypesafeDatastoreConfigAdapter.java | 2 ++ .../model/config/ConnectionConfig.java | 2 ++ .../postgres/PostgresConnectionConfig.java | 21 +++++++++++++++++++ .../postgres/PostgresDatastore.java | 6 +++++- .../postgres/PostgresSchemaRegistry.java | 17 --------------- 5 files changed, 30 insertions(+), 18 deletions(-) diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/TypesafeDatastoreConfigAdapter.java b/document-store/src/main/java/org/hypertrace/core/documentstore/TypesafeDatastoreConfigAdapter.java index c1fcacde..dda12fff 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/TypesafeDatastoreConfigAdapter.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/TypesafeDatastoreConfigAdapter.java @@ -93,6 +93,8 @@ public DatastoreConfig convert(final Config config) { connectionConfig.credentials(), connectionConfig.applicationName(), connectionConfig.connectionPoolConfig(), + connectionConfig.schemaCacheExpiry(), + connectionConfig.schemaRefreshCooldown(), connectionConfig.customParameters()) { @Override public String toConnectionString() { diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/model/config/ConnectionConfig.java b/document-store/src/main/java/org/hypertrace/core/documentstore/model/config/ConnectionConfig.java index 18d9975c..2960a2d2 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/model/config/ConnectionConfig.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/model/config/ConnectionConfig.java @@ -126,6 +126,8 @@ public ConnectionConfig build() { credentials, applicationName, connectionPoolConfig, + null, + null, customParameters); } diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/model/config/postgres/PostgresConnectionConfig.java b/document-store/src/main/java/org/hypertrace/core/documentstore/model/config/postgres/PostgresConnectionConfig.java index 925e4b1e..2d1ab176 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/model/config/postgres/PostgresConnectionConfig.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/model/config/postgres/PostgresConnectionConfig.java @@ -3,6 +3,7 @@ import static java.util.Collections.unmodifiableList; import static java.util.function.Predicate.not; +import java.time.Duration; import java.util.ArrayList; import java.util.List; import java.util.Map; @@ -34,8 +35,13 @@ public class PostgresConnectionConfig extends ConnectionConfig { .password(PostgresDefaults.DEFAULT_PASSWORD) .build(); + private static final Duration DEFAULT_SCHEMA_CACHE_EXPIRY = Duration.ofHours(24); + private static final Duration DEFAULT_SCHEMA_REFRESH_COOLDOWN = Duration.ofMinutes(15); + @NonNull String applicationName; @NonNull ConnectionPoolConfig connectionPoolConfig; + @NonNull Duration schemaCacheExpiry; + @NonNull Duration schemaRefreshCooldown; public static ConnectionConfigBuilder builder() { return ConnectionConfig.builder().type(DatabaseType.POSTGRES); @@ -47,6 +53,8 @@ public PostgresConnectionConfig( @Nullable final ConnectionCredentials credentials, @NonNull final String applicationName, @Nullable final ConnectionPoolConfig connectionPoolConfig, + @Nullable final Duration schemaCacheExpiry, + @Nullable final Duration schemaRefreshCooldown, @NonNull final Map customParameters) { super( ensureSingleEndpoint(endpoints), @@ -55,6 +63,8 @@ public PostgresConnectionConfig( customParameters); this.applicationName = applicationName; this.connectionPoolConfig = getConnectionPoolConfigOrDefault(connectionPoolConfig); + this.schemaCacheExpiry = getSchemaCacheExpiryOrDefault(schemaCacheExpiry); + this.schemaRefreshCooldown = getSchemaRefreshCooldownOrDefault(schemaRefreshCooldown); } public String toConnectionString() { @@ -127,4 +137,15 @@ private ConnectionPoolConfig getConnectionPoolConfigOrDefault( @Nullable final ConnectionPoolConfig connectionPoolConfig) { return Optional.ofNullable(connectionPoolConfig).orElse(ConnectionPoolConfig.builder().build()); } + + @NonNull + private Duration getSchemaCacheExpiryOrDefault(@Nullable final Duration schemaCacheExpiry) { + return Optional.ofNullable(schemaCacheExpiry).orElse(DEFAULT_SCHEMA_CACHE_EXPIRY); + } + + @NonNull + private Duration getSchemaRefreshCooldownOrDefault( + @Nullable final Duration schemaRefreshCooldown) { + return Optional.ofNullable(schemaRefreshCooldown).orElse(DEFAULT_SCHEMA_REFRESH_COOLDOWN); + } } 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 3b1fb345..4e9fb3e2 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 @@ -60,7 +60,11 @@ public PostgresDatastore(@NonNull final DatastoreConfig datastoreConfig) { client = new PostgresClient(postgresConnectionConfig); database = connectionConfig.database(); docStoreMetricProvider = new PostgresDocStoreMetricProvider(this, postgresConnectionConfig); - schemaRegistry = new PostgresSchemaRegistry(new PostgresMetadataFetcher(client)); + schemaRegistry = + new PostgresSchemaRegistry( + new PostgresMetadataFetcher(client), + postgresConnectionConfig.schemaCacheExpiry(), + postgresConnectionConfig.schemaRefreshCooldown()); } catch (final IllegalArgumentException e) { throw new IllegalArgumentException( String.format("Unable to instantiate PostgresClient with config:%s", connectionConfig), diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistry.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistry.java index ff7d6164..49f3a0c2 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistry.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistry.java @@ -52,28 +52,11 @@ */ public class PostgresSchemaRegistry implements SchemaRegistry { - /** Default cache expiry time: 24 hours. */ - private static final Duration DEFAULT_CACHE_EXPIRY = Duration.ofHours(24); - - /** Default cooldown period between refresh attempts: 15 minutes. */ - private static final Duration DEFAULT_REFRESH_COOLDOWN = Duration.ofMinutes(15); - private final LoadingCache> cache; private final Map lastRefreshTimes; private final Duration refreshCooldown; private final Clock clock; - /** - * Creates a new schema registry with default settings. - * - *

Uses default cache expiry of 24 hours and refresh cooldown of 15 minutes. - * - * @param fetcher the metadata fetcher to use for loading schema information - */ - public PostgresSchemaRegistry(PostgresMetadataFetcher fetcher) { - this(fetcher, DEFAULT_CACHE_EXPIRY, DEFAULT_REFRESH_COOLDOWN, Clock.systemUTC()); - } - /** * Creates a new schema registry with custom cache settings. * From 7bf77c5cf7b39ac3cfe04086f57b53c63f618fcb Mon Sep 17 00:00:00 2001 From: Prashant Pandey Date: Mon, 29 Dec 2025 15:55:12 +0530 Subject: [PATCH 10/26] Added PostgresMetadataFetcherTest --- .../expression/impl/DataType.java | 1 + .../postgres/PostgresMetadataFetcher.java | 4 + .../nonjson/field/PostgresDataType.java | 1 + .../postgres/PostgresMetadataFetcherTest.java | 310 ++++++++++++++++++ 4 files changed, 316 insertions(+) create mode 100644 document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresMetadataFetcherTest.java diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/expression/impl/DataType.java b/document-store/src/main/java/org/hypertrace/core/documentstore/expression/impl/DataType.java index 9eec68b9..44ed4d2c 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/expression/impl/DataType.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/expression/impl/DataType.java @@ -22,6 +22,7 @@ public enum DataType { FLOAT, DOUBLE, BOOLEAN, + JSON, // timestamp with time-zone information. For example: 2004-10-19 10:23:54+02. // For more info, see: https://www.postgresql.org/docs/current/datatype-datetime.html TIMESTAMPTZ, diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresMetadataFetcher.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresMetadataFetcher.java index e61b26b7..2d2be0ec 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresMetadataFetcher.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresMetadataFetcher.java @@ -81,6 +81,8 @@ private DataType mapToCanonicalType(String udtName) { case "bpchar": case "uuid": return DataType.STRING; + case "jsonb": + return DataType.JSON; default: return DataType.UNSPECIFIED; } @@ -114,6 +116,8 @@ private PostgresDataType mapToPostgresType(String udtName) { case "bpchar": case "uuid": return PostgresDataType.TEXT; + case "jsonb": + return PostgresDataType.JSONB; default: return PostgresDataType.UNKNOWN; } diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/parser/filter/nonjson/field/PostgresDataType.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/parser/filter/nonjson/field/PostgresDataType.java index 89626e7c..c920473f 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/parser/filter/nonjson/field/PostgresDataType.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/parser/filter/nonjson/field/PostgresDataType.java @@ -15,6 +15,7 @@ public enum PostgresDataType { REAL("float4"), DOUBLE_PRECISION("float8"), BOOLEAN("bool"), + JSONB("jsonb"), TIMESTAMPTZ("timestamptz"), DATE("date"), UNKNOWN(null); diff --git a/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresMetadataFetcherTest.java b/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresMetadataFetcherTest.java new file mode 100644 index 00000000..6fa32d91 --- /dev/null +++ b/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresMetadataFetcherTest.java @@ -0,0 +1,310 @@ +package org.hypertrace.core.documentstore.postgres; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.util.Map; +import org.hypertrace.core.documentstore.expression.impl.DataType; +import org.hypertrace.core.documentstore.postgres.model.PostgresColumnMetadata; +import org.hypertrace.core.documentstore.postgres.query.v1.parser.filter.nonjson.field.PostgresDataType; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; + +@ExtendWith(MockitoExtension.class) +class PostgresMetadataFetcherTest { + + private static final String TEST_TABLE = "test_table"; + + @Mock private PostgresClient client; + @Mock private Connection connection; + @Mock private PreparedStatement preparedStatement; + @Mock private ResultSet resultSet; + + private PostgresMetadataFetcher fetcher; + + @BeforeEach + void setUp() throws SQLException { + when(client.getPooledConnection()).thenReturn(connection); + when(connection.prepareStatement(anyString())).thenReturn(preparedStatement); + when(preparedStatement.executeQuery()).thenReturn(resultSet); + fetcher = new PostgresMetadataFetcher(client); + } + + @Test + void fetchReturnsEmptyMapForTableWithNoColumns() throws SQLException { + when(resultSet.next()).thenReturn(false); + + Map result = fetcher.fetch(TEST_TABLE); + + assertTrue(result.isEmpty()); + verify(preparedStatement).setString(1, TEST_TABLE); + } + + @Test + void fetchReturnsSingleColumn() throws SQLException { + when(resultSet.next()).thenReturn(true, false); + when(resultSet.getString("column_name")).thenReturn("id"); + when(resultSet.getString("udt_name")).thenReturn("int4"); + when(resultSet.getString("is_nullable")).thenReturn("NO"); + + Map result = fetcher.fetch(TEST_TABLE); + + assertEquals(1, result.size()); + PostgresColumnMetadata metadata = result.get("id"); + assertNotNull(metadata); + assertEquals("id", metadata.getName()); + assertEquals(DataType.INTEGER, metadata.getCanonicalType()); + assertEquals(PostgresDataType.INTEGER, metadata.getPostgresType()); + assertFalse(metadata.isNullable()); + } + + @Test + void fetchReturnsMultipleColumns() throws SQLException { + when(resultSet.next()).thenReturn(true, true, true, false); + when(resultSet.getString("column_name")).thenReturn("id", "name", "price"); + when(resultSet.getString("udt_name")).thenReturn("int8", "text", "float8"); + when(resultSet.getString("is_nullable")).thenReturn("NO", "YES", "YES"); + + Map result = fetcher.fetch(TEST_TABLE); + + assertEquals(3, result.size()); + + // Verify id column + PostgresColumnMetadata idMeta = result.get("id"); + assertEquals(DataType.LONG, idMeta.getCanonicalType()); + assertEquals(PostgresDataType.BIGINT, idMeta.getPostgresType()); + assertFalse(idMeta.isNullable()); + + // Verify name column + PostgresColumnMetadata nameMeta = result.get("name"); + assertEquals(DataType.STRING, nameMeta.getCanonicalType()); + assertEquals(PostgresDataType.TEXT, nameMeta.getPostgresType()); + assertTrue(nameMeta.isNullable()); + + // Verify price column + PostgresColumnMetadata priceMeta = result.get("price"); + assertEquals(DataType.DOUBLE, priceMeta.getCanonicalType()); + assertEquals(PostgresDataType.DOUBLE_PRECISION, priceMeta.getPostgresType()); + assertTrue(priceMeta.isNullable()); + } + + @Test + void fetchMapsInt4ToInteger() throws SQLException { + setupSingleColumnResult("col", "int4", "NO"); + + PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col"); + + assertEquals(DataType.INTEGER, meta.getCanonicalType()); + assertEquals(PostgresDataType.INTEGER, meta.getPostgresType()); + } + + @Test + void fetchMapsInt2ToInteger() throws SQLException { + setupSingleColumnResult("col", "int2", "NO"); + + PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col"); + + assertEquals(DataType.INTEGER, meta.getCanonicalType()); + assertEquals(PostgresDataType.INTEGER, meta.getPostgresType()); + } + + @Test + void fetchMapsInt8ToLong() throws SQLException { + setupSingleColumnResult("col", "int8", "NO"); + + PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col"); + + assertEquals(DataType.LONG, meta.getCanonicalType()); + assertEquals(PostgresDataType.BIGINT, meta.getPostgresType()); + } + + @Test + void fetchMapsFloat4ToFloat() throws SQLException { + setupSingleColumnResult("col", "float4", "NO"); + + PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col"); + + assertEquals(DataType.FLOAT, meta.getCanonicalType()); + assertEquals(PostgresDataType.REAL, meta.getPostgresType()); + } + + @Test + void fetchMapsFloat8ToDouble() throws SQLException { + setupSingleColumnResult("col", "float8", "NO"); + + PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col"); + + assertEquals(DataType.DOUBLE, meta.getCanonicalType()); + assertEquals(PostgresDataType.DOUBLE_PRECISION, meta.getPostgresType()); + } + + @Test + void fetchMapsNumericToDouble() throws SQLException { + setupSingleColumnResult("col", "numeric", "NO"); + + PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col"); + + assertEquals(DataType.DOUBLE, meta.getCanonicalType()); + assertEquals(PostgresDataType.DOUBLE_PRECISION, meta.getPostgresType()); + } + + @Test + void fetchMapsBoolToBoolean() throws SQLException { + setupSingleColumnResult("col", "bool", "NO"); + + PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col"); + + assertEquals(DataType.BOOLEAN, meta.getCanonicalType()); + assertEquals(PostgresDataType.BOOLEAN, meta.getPostgresType()); + } + + @Test + void fetchMapsTextToString() throws SQLException { + setupSingleColumnResult("col", "text", "NO"); + + PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col"); + + assertEquals(DataType.STRING, meta.getCanonicalType()); + assertEquals(PostgresDataType.TEXT, meta.getPostgresType()); + } + + @Test + void fetchMapsVarcharToString() throws SQLException { + setupSingleColumnResult("col", "varchar", "NO"); + + PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col"); + + assertEquals(DataType.STRING, meta.getCanonicalType()); + assertEquals(PostgresDataType.TEXT, meta.getPostgresType()); + } + + @Test + void fetchMapsBpcharToString() throws SQLException { + setupSingleColumnResult("col", "bpchar", "NO"); + + PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col"); + + assertEquals(DataType.STRING, meta.getCanonicalType()); + assertEquals(PostgresDataType.TEXT, meta.getPostgresType()); + } + + @Test + void fetchMapsUuidToString() throws SQLException { + setupSingleColumnResult("col", "uuid", "NO"); + + PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col"); + + assertEquals(DataType.STRING, meta.getCanonicalType()); + assertEquals(PostgresDataType.TEXT, meta.getPostgresType()); + } + + @Test + void fetchMapsJsonbToJson() throws SQLException { + setupSingleColumnResult("col", "jsonb", "NO"); + + PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col"); + + assertEquals(DataType.JSON, meta.getCanonicalType()); + assertEquals(PostgresDataType.JSONB, meta.getPostgresType()); + } + + @Test + void fetchMapsTimestamptzToTimestamptz() throws SQLException { + setupSingleColumnResult("col", "timestamptz", "NO"); + + PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col"); + + assertEquals(DataType.TIMESTAMPTZ, meta.getCanonicalType()); + assertEquals(PostgresDataType.TIMESTAMPTZ, meta.getPostgresType()); + } + + @Test + void fetchMapsDateToDate() throws SQLException { + setupSingleColumnResult("col", "date", "NO"); + + PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col"); + + assertEquals(DataType.DATE, meta.getCanonicalType()); + assertEquals(PostgresDataType.DATE, meta.getPostgresType()); + } + + @Test + void fetchMapsUnknownTypeToUnspecified() throws SQLException { + setupSingleColumnResult("col", "unknown_type", "NO"); + + PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col"); + + assertEquals(DataType.UNSPECIFIED, meta.getCanonicalType()); + assertEquals(PostgresDataType.UNKNOWN, meta.getPostgresType()); + } + + @Test + void fetchMapsNullUdtNameToUnspecified() throws SQLException { + setupSingleColumnResult("col", null, "NO"); + + PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col"); + + assertEquals(DataType.UNSPECIFIED, meta.getCanonicalType()); + assertEquals(PostgresDataType.UNKNOWN, meta.getPostgresType()); + } + + @Test + void fetchHandlesNullableColumn() throws SQLException { + setupSingleColumnResult("col", "text", "YES"); + + PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col"); + + assertTrue(meta.isNullable()); + } + + @Test + void fetchHandlesNonNullableColumn() throws SQLException { + setupSingleColumnResult("col", "text", "NO"); + + PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col"); + + assertFalse(meta.isNullable()); + } + + @Test + void fetchHandlesCaseInsensitiveUdtName() throws SQLException { + setupSingleColumnResult("col", "INT4", "NO"); + + PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col"); + + assertEquals(DataType.INTEGER, meta.getCanonicalType()); + assertEquals(PostgresDataType.INTEGER, meta.getPostgresType()); + } + + @Test + void fetchThrowsRuntimeExceptionOnSqlException() throws SQLException { + when(preparedStatement.executeQuery()).thenThrow(new SQLException("Connection failed")); + + RuntimeException exception = + assertThrows(RuntimeException.class, () -> fetcher.fetch(TEST_TABLE)); + + assertTrue(exception.getMessage().contains(TEST_TABLE)); + assertTrue(exception.getCause() instanceof SQLException); + } + + private void setupSingleColumnResult(String colName, String udtName, String isNullable) + throws SQLException { + when(resultSet.next()).thenReturn(true, false); + when(resultSet.getString("column_name")).thenReturn(colName); + when(resultSet.getString("udt_name")).thenReturn(udtName); + when(resultSet.getString("is_nullable")).thenReturn(isNullable); + } +} From 6d03cd53c77c2c598da9d05e0493b7fb8258002f Mon Sep 17 00:00:00 2001 From: Prashant Pandey Date: Mon, 29 Dec 2025 16:04:20 +0530 Subject: [PATCH 11/26] WIP --- .../postgres/PostgresSchemaRegistry.java | 23 +------- .../postgres/PostgresSchemaRegistryTest.java | 57 ++++--------------- 2 files changed, 13 insertions(+), 67 deletions(-) diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistry.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistry.java index 49f3a0c2..c4a0ae98 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistry.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistry.java @@ -3,7 +3,6 @@ import com.google.common.cache.CacheBuilder; import com.google.common.cache.CacheLoader; import com.google.common.cache.LoadingCache; -import java.time.Clock; import java.time.Duration; import java.time.Instant; import java.util.Map; @@ -55,7 +54,6 @@ public class PostgresSchemaRegistry implements SchemaRegistry> cache; private final Map lastRefreshTimes; private final Duration refreshCooldown; - private final Clock clock; /** * Creates a new schema registry with custom cache settings. @@ -66,24 +64,7 @@ public class PostgresSchemaRegistry implements SchemaRegistry(); this.cache = CacheBuilder.newBuilder() @@ -92,7 +73,7 @@ public PostgresSchemaRegistry( new CacheLoader<>() { @Override public Map load(String tableName) { - lastRefreshTimes.put(tableName, clock.instant()); + lastRefreshTimes.put(tableName, Instant.now()); return fetcher.fetch(tableName); } }); @@ -167,6 +148,6 @@ private boolean canRefresh(String tableName) { if (lastRefresh == null) { return true; } - return Duration.between(lastRefresh, clock.instant()).compareTo(refreshCooldown) >= 0; + return Duration.between(lastRefresh, Instant.now()).compareTo(refreshCooldown) >= 0; } } diff --git a/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistryTest.java b/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistryTest.java index fa8e4a03..7fa22d5f 100644 --- a/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistryTest.java +++ b/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistryTest.java @@ -8,10 +8,7 @@ import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; -import java.time.Clock; import java.time.Duration; -import java.time.Instant; -import java.time.ZoneId; import java.util.HashMap; import java.util.Map; import java.util.Optional; @@ -32,17 +29,15 @@ class PostgresSchemaRegistryTest { private static final String COL_NAME = "name"; private static final String COL_PRICE = "price"; private static final Duration CACHE_EXPIRY = Duration.ofHours(24); - private static final Duration REFRESH_COOLDOWN = Duration.ofMinutes(15); + private static final Duration REFRESH_COOLDOWN = Duration.ofMillis(50); @Mock private PostgresMetadataFetcher fetcher; private PostgresSchemaRegistry registry; - private MutableClock mutableClock; @BeforeEach void setUp() { - mutableClock = new MutableClock(Instant.parse("2024-01-01T00:00:00Z")); - registry = new PostgresSchemaRegistry(fetcher, CACHE_EXPIRY, REFRESH_COOLDOWN, mutableClock); + registry = new PostgresSchemaRegistry(fetcher, CACHE_EXPIRY, REFRESH_COOLDOWN); } @Test @@ -157,7 +152,7 @@ void getColumnOrRefreshReturnsColumnIfExists() { } @Test - void getColumnOrRefreshRefreshesSchemaIfColumnMissingAndCooldownExpired() { + void getColumnOrRefreshRefreshesSchemaIfColumnMissingAndCooldownExpired() throws Exception { // Initial schema without the "new_col" Map initialSchema = createTestSchema(); @@ -179,8 +174,8 @@ void getColumnOrRefreshRefreshesSchemaIfColumnMissingAndCooldownExpired() { registry.getSchema(TEST_TABLE); verify(fetcher, times(1)).fetch(TEST_TABLE); - // Advance time past cooldown period - mutableClock.advance(REFRESH_COOLDOWN.plusMinutes(1)); + // Wait past cooldown period + Thread.sleep(REFRESH_COOLDOWN.toMillis() + 10); // Now try to get missing column - should trigger refresh Optional result = registry.getColumnOrRefresh(TEST_TABLE, "new_col"); @@ -210,7 +205,7 @@ void getColumnOrRefreshDoesNotRefreshIfWithinCooldownPeriod() { } @Test - void getColumnOrRefreshRefreshesAfterCooldownExpires() { + void getColumnOrRefreshRefreshesAfterCooldownExpires() throws Exception { Map schema = createTestSchema(); when(fetcher.fetch(TEST_TABLE)).thenReturn(schema); @@ -222,8 +217,8 @@ void getColumnOrRefreshRefreshesAfterCooldownExpires() { registry.getColumnOrRefresh(TEST_TABLE, "nonexistent_col"); verify(fetcher, times(1)).fetch(TEST_TABLE); - // Advance time past cooldown - mutableClock.advance(REFRESH_COOLDOWN.plusSeconds(1)); + // Wait past cooldown + Thread.sleep(REFRESH_COOLDOWN.toMillis() + 10); // Try again - should refresh now registry.getColumnOrRefresh(TEST_TABLE, "nonexistent_col"); @@ -231,15 +226,15 @@ void getColumnOrRefreshRefreshesAfterCooldownExpires() { } @Test - void getColumnOrRefreshReturnsNullIfColumnStillMissingAfterRefresh() { + void getColumnOrRefreshReturnsEmptyIfColumnStillMissingAfterRefresh() throws Exception { Map schema = createTestSchema(); when(fetcher.fetch(TEST_TABLE)).thenReturn(schema); // First call loads the schema registry.getSchema(TEST_TABLE); - // Advance past cooldown - mutableClock.advance(REFRESH_COOLDOWN.plusMinutes(1)); + // Wait past cooldown + Thread.sleep(REFRESH_COOLDOWN.toMillis() + 10); // Try to get a column that doesn't exist even after refresh Optional result = @@ -308,34 +303,4 @@ private Map createTestSchema() { .build()); return schema; } - - /** A mutable clock for testing time-dependent behavior. */ - private static class MutableClock extends Clock { - private Instant currentInstant; - private final ZoneId zone; - - MutableClock(Instant initialInstant) { - this.currentInstant = initialInstant; - this.zone = ZoneId.of("UTC"); - } - - void advance(Duration duration) { - currentInstant = currentInstant.plus(duration); - } - - @Override - public ZoneId getZone() { - return zone; - } - - @Override - public Clock withZone(ZoneId zone) { - return this; - } - - @Override - public Instant instant() { - return currentInstant; - } - } } From c3f5f7e9ba7ff93a51384d8d23b658dce7ec4c1c Mon Sep 17 00:00:00 2001 From: Prashant Pandey Date: Mon, 29 Dec 2025 16:27:49 +0530 Subject: [PATCH 12/26] Added docs on thread safety --- .../documentstore/postgres/PostgresSchemaRegistry.java | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistry.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistry.java index c4a0ae98..e8fc9b19 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistry.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistry.java @@ -52,7 +52,7 @@ public class PostgresSchemaRegistry implements SchemaRegistry { private final LoadingCache> cache; - private final Map lastRefreshTimes; + private final ConcurrentHashMap lastRefreshTimes; private final Duration refreshCooldown; /** @@ -121,6 +121,11 @@ public void invalidate(String tableName) { * returned without hitting the database. * * + *

Note that this is a check-then-act sequence that should ideally be atomic. However, this + * method is deliberately not thread-safe since even in case of a data race, it will result in one + * extra call to the DB, which will not be allowed anyway due to the cooldown period having been + * reset by the previous call. This is likely to be more performant than contending for locks. + * * @param tableName the name of the table * @param colName the name of the column * @return an Optional containing the column metadata, or empty if the column does not exist From 827381f7a9111e666f0c4a7ff2ed465f5a997c73 Mon Sep 17 00:00:00 2001 From: Prashant Pandey Date: Mon, 29 Dec 2025 16:34:54 +0530 Subject: [PATCH 13/26] Added PostgresSchemaRegistryIntegrationTest.java --- ...PostgresSchemaRegistryIntegrationTest.java | 279 ++++++++++++++++++ 1 file changed, 279 insertions(+) create mode 100644 document-store/src/integrationTest/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistryIntegrationTest.java diff --git a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistryIntegrationTest.java b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistryIntegrationTest.java new file mode 100644 index 00000000..9286d7e5 --- /dev/null +++ b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistryIntegrationTest.java @@ -0,0 +1,279 @@ +package org.hypertrace.core.documentstore.postgres; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.typesafe.config.Config; +import com.typesafe.config.ConfigFactory; +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.util.HashMap; +import java.util.Map; +import java.util.Optional; +import org.hypertrace.core.documentstore.DatastoreProvider; +import org.hypertrace.core.documentstore.commons.SchemaRegistry; +import org.hypertrace.core.documentstore.expression.impl.DataType; +import org.hypertrace.core.documentstore.postgres.model.PostgresColumnMetadata; +import org.hypertrace.core.documentstore.postgres.query.v1.parser.filter.nonjson.field.PostgresDataType; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.testcontainers.containers.GenericContainer; +import org.testcontainers.containers.wait.strategy.Wait; +import org.testcontainers.junit.jupiter.Testcontainers; +import org.testcontainers.utility.DockerImageName; + +@Testcontainers +class PostgresSchemaRegistryIntegrationTest { + + private static final String TABLE_NAME = "myTestFlat"; + + private static GenericContainer postgres; + private static PostgresDatastore datastore; + private static SchemaRegistry registry; + + @BeforeAll + static void init() throws Exception { + postgres = + new GenericContainer<>(DockerImageName.parse("postgres:13.1")) + .withEnv("POSTGRES_PASSWORD", "postgres") + .withEnv("POSTGRES_USER", "postgres") + .withExposedPorts(5432) + .waitingFor(Wait.forListeningPort()); + postgres.start(); + + String connectionUrl = + String.format("jdbc:postgresql://localhost:%s/", postgres.getMappedPort(5432)); + + Map postgresConfig = new HashMap<>(); + postgresConfig.put("url", connectionUrl); + postgresConfig.put("user", "postgres"); + postgresConfig.put("password", "postgres"); + Config config = ConfigFactory.parseMap(postgresConfig); + + datastore = (PostgresDatastore) DatastoreProvider.getDatastore("Postgres", config); + + createFlatTable(); + + registry = datastore.getSchemaRegistry(); + } + + private static void createFlatTable() throws Exception { + String createTableSQL = + String.format( + "CREATE TABLE IF NOT EXISTS \"%s\" (" + + "\"_id\" INTEGER PRIMARY KEY," + + "\"item\" TEXT," + + "\"price\" INTEGER," + + "\"quantity\" BIGINT," + + "\"rating\" REAL," + + "\"score\" DOUBLE PRECISION," + + "\"date\" TIMESTAMPTZ," + + "\"created_date\" DATE," + + "\"in_stock\" BOOLEAN," + + "\"tags\" TEXT[]," + + "\"props\" JSONB" + + ");", + TABLE_NAME); + + try (Connection connection = datastore.getPostgresClient(); + PreparedStatement statement = connection.prepareStatement(createTableSQL)) { + statement.execute(); + System.out.println("Created flat table: " + TABLE_NAME); + } + } + + @BeforeEach + void setUp() { + registry.invalidate(TABLE_NAME); + } + + @AfterAll + static void shutdown() { + if (postgres != null) { + postgres.stop(); + } + } + + @Test + void getSchemaReturnsAllColumns() { + Map schema = registry.getSchema(TABLE_NAME); + + assertNotNull(schema); + assertEquals(11, schema.size()); + assertTrue(schema.containsKey("_id")); + assertTrue(schema.containsKey("item")); + assertTrue(schema.containsKey("price")); + assertTrue(schema.containsKey("quantity")); + assertTrue(schema.containsKey("rating")); + assertTrue(schema.containsKey("score")); + assertTrue(schema.containsKey("date")); + assertTrue(schema.containsKey("created_date")); + assertTrue(schema.containsKey("in_stock")); + assertTrue(schema.containsKey("tags")); + assertTrue(schema.containsKey("props")); + } + + @Test + void getSchemaReturnsCorrectIntegerMapping() { + Map schema = registry.getSchema(TABLE_NAME); + + PostgresColumnMetadata idMeta = schema.get("_id"); + assertEquals("_id", idMeta.getName()); + assertEquals(DataType.INTEGER, idMeta.getCanonicalType()); + assertEquals(PostgresDataType.INTEGER, idMeta.getPostgresType()); + assertFalse(idMeta.isNullable()); + + PostgresColumnMetadata priceMeta = schema.get("price"); + assertEquals(DataType.INTEGER, priceMeta.getCanonicalType()); + assertEquals(PostgresDataType.INTEGER, priceMeta.getPostgresType()); + assertTrue(priceMeta.isNullable()); + } + + @Test + void getSchemaReturnsCorrectBigintMapping() { + Map schema = registry.getSchema(TABLE_NAME); + + PostgresColumnMetadata quantityMeta = schema.get("quantity"); + assertEquals(DataType.LONG, quantityMeta.getCanonicalType()); + assertEquals(PostgresDataType.BIGINT, quantityMeta.getPostgresType()); + } + + @Test + void getSchemaReturnsCorrectFloatMapping() { + Map schema = registry.getSchema(TABLE_NAME); + + PostgresColumnMetadata ratingMeta = schema.get("rating"); + assertEquals(DataType.FLOAT, ratingMeta.getCanonicalType()); + assertEquals(PostgresDataType.REAL, ratingMeta.getPostgresType()); + } + + @Test + void getSchemaReturnsCorrectDoubleMapping() { + Map schema = registry.getSchema(TABLE_NAME); + + PostgresColumnMetadata scoreMeta = schema.get("score"); + assertEquals(DataType.DOUBLE, scoreMeta.getCanonicalType()); + assertEquals(PostgresDataType.DOUBLE_PRECISION, scoreMeta.getPostgresType()); + } + + @Test + void getSchemaReturnsCorrectTextMapping() { + Map schema = registry.getSchema(TABLE_NAME); + + PostgresColumnMetadata itemMeta = schema.get("item"); + assertEquals(DataType.STRING, itemMeta.getCanonicalType()); + assertEquals(PostgresDataType.TEXT, itemMeta.getPostgresType()); + } + + @Test + void getSchemaReturnsCorrectBooleanMapping() { + Map schema = registry.getSchema(TABLE_NAME); + + PostgresColumnMetadata inStockMeta = schema.get("in_stock"); + assertEquals(DataType.BOOLEAN, inStockMeta.getCanonicalType()); + assertEquals(PostgresDataType.BOOLEAN, inStockMeta.getPostgresType()); + } + + @Test + void getSchemaReturnsCorrectTimestamptzMapping() { + Map schema = registry.getSchema(TABLE_NAME); + + PostgresColumnMetadata dateMeta = schema.get("date"); + assertEquals(DataType.TIMESTAMPTZ, dateMeta.getCanonicalType()); + assertEquals(PostgresDataType.TIMESTAMPTZ, dateMeta.getPostgresType()); + } + + @Test + void getSchemaReturnsCorrectDateMapping() { + Map schema = registry.getSchema(TABLE_NAME); + + PostgresColumnMetadata createdDateMeta = schema.get("created_date"); + assertEquals(DataType.DATE, createdDateMeta.getCanonicalType()); + assertEquals(PostgresDataType.DATE, createdDateMeta.getPostgresType()); + } + + @Test + void getSchemaReturnsCorrectJsonbMapping() { + Map schema = registry.getSchema(TABLE_NAME); + + PostgresColumnMetadata propsMeta = schema.get("props"); + assertEquals(DataType.JSON, propsMeta.getCanonicalType()); + assertEquals(PostgresDataType.JSONB, propsMeta.getPostgresType()); + } + + @Test + void getColumnOrRefreshReturnsExistingColumn() { + Optional result = registry.getColumnOrRefresh(TABLE_NAME, "item"); + + assertTrue(result.isPresent()); + assertEquals("item", result.get().getName()); + assertEquals(DataType.STRING, result.get().getCanonicalType()); + } + + @Test + void getColumnOrRefreshReturnsEmptyForNonExistentColumn() { + Optional result = + registry.getColumnOrRefresh(TABLE_NAME, "nonexistent_column"); + + assertFalse(result.isPresent()); + } + + @Test + void getColumnOrRefreshFindsNewlyAddedColumnAfterInvalidation() throws Exception { + // First, verify the new column doesn't exist + Optional before = registry.getColumnOrRefresh(TABLE_NAME, "new_column"); + assertFalse(before.isPresent()); + + // Add a new column to the table + try (Connection connection = datastore.getPostgresClient(); + PreparedStatement statement = + connection.prepareStatement( + String.format("ALTER TABLE \"%s\" ADD COLUMN \"new_column\" TEXT", TABLE_NAME))) { + statement.execute(); + } + + // Invalidate cache to force reload + registry.invalidate(TABLE_NAME); + + // Now the registry should find the new column after reload + Optional after = registry.getColumnOrRefresh(TABLE_NAME, "new_column"); + assertTrue(after.isPresent()); + assertEquals("new_column", after.get().getName()); + assertEquals(DataType.STRING, after.get().getCanonicalType()); + + // Cleanup: drop the column + try (Connection connection = datastore.getPostgresClient(); + PreparedStatement statement = + connection.prepareStatement( + String.format("ALTER TABLE \"%s\" DROP COLUMN \"new_column\"", TABLE_NAME))) { + statement.execute(); + } + } + + @Test + void cacheReturnsSameInstanceOnSubsequentCalls() { + Map schema1 = registry.getSchema(TABLE_NAME); + Map schema2 = registry.getSchema(TABLE_NAME); + + // Should be the same cached instance + assertTrue(schema1 == schema2); + } + + @Test + void invalidateCausesReload() { + Map schema1 = registry.getSchema(TABLE_NAME); + + registry.invalidate(TABLE_NAME); + + Map schema2 = registry.getSchema(TABLE_NAME); + + // Should be different instances after invalidation + assertFalse(schema1 == schema2); + // But same content + assertEquals(schema1.keySet(), schema2.keySet()); + } +} From 602037bfb07552a3bd17af7414606892f84bb032 Mon Sep 17 00:00:00 2001 From: Prashant Pandey Date: Mon, 29 Dec 2025 16:36:28 +0530 Subject: [PATCH 14/26] WIP --- .../postgres/PostgresSchemaRegistryIntegrationTest.java | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistryIntegrationTest.java b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistryIntegrationTest.java index 9286d7e5..2b3b30d4 100644 --- a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistryIntegrationTest.java +++ b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistryIntegrationTest.java @@ -3,6 +3,8 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNotSame; +import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertTrue; import com.typesafe.config.Config; @@ -260,7 +262,7 @@ void cacheReturnsSameInstanceOnSubsequentCalls() { Map schema2 = registry.getSchema(TABLE_NAME); // Should be the same cached instance - assertTrue(schema1 == schema2); + assertSame(schema1, schema2); } @Test @@ -272,7 +274,7 @@ void invalidateCausesReload() { Map schema2 = registry.getSchema(TABLE_NAME); // Should be different instances after invalidation - assertFalse(schema1 == schema2); + assertNotSame(schema1, schema2); // But same content assertEquals(schema1.keySet(), schema2.keySet()); } From c8a53eb3bf42c618f0ead1c964dea91571525b92 Mon Sep 17 00:00:00 2001 From: Prashant Pandey Date: Tue, 30 Dec 2025 10:59:40 +0530 Subject: [PATCH 15/26] WIP --- .../FlatCollectionWriteTest.java | 105 +++++++++- .../postgres/FlatPostgresCollection.java | 186 +++++++++++++++++- 2 files changed, 282 insertions(+), 9 deletions(-) diff --git a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java index 1595799c..65d26904 100644 --- a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java +++ b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java @@ -1,7 +1,9 @@ package org.hypertrace.core.documentstore; import static org.hypertrace.core.documentstore.utils.Utils.readFileFromResource; +import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.ObjectMapper; @@ -10,10 +12,12 @@ import java.io.IOException; import java.sql.Connection; import java.sql.PreparedStatement; +import java.sql.ResultSet; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Set; +import org.hypertrace.core.documentstore.model.exception.DuplicateDocumentException; import org.hypertrace.core.documentstore.postgres.PostgresDatastore; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.AfterEach; @@ -200,22 +204,109 @@ void testUpsertAndReturn() { class CreateTests { @Test - @DisplayName("Should throw UnsupportedOperationException for create") - void testCreate() { + @DisplayName("Should create a new document with all field types") + void testCreateNewDocument() throws Exception { ObjectNode objectNode = OBJECT_MAPPER.createObjectNode(); - objectNode.put("_id", 300); + objectNode.put("id", "new-doc-100"); objectNode.put("item", "Brand New Item"); + objectNode.put("price", 999); + objectNode.put("quantity", 50); + objectNode.put("in_stock", true); Document document = new JSONDocument(objectNode); - Key key = new SingleValueKey("default", "300"); + Key key = new SingleValueKey("default", "new-doc-100"); + + CreateResult result = flatCollection.create(key, document); + + assertTrue(result.isSucceed()); + + // Verify the data was inserted + PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; + try (Connection conn = pgDatastore.getPostgresClient(); + PreparedStatement ps = + conn.prepareStatement( + String.format( + "SELECT * FROM \"%s\" WHERE \"id\" = 'new-doc-100'", FLAT_COLLECTION_NAME)); + ResultSet rs = ps.executeQuery()) { + assertTrue(rs.next()); + assertEquals("Brand New Item", rs.getString("item")); + assertEquals(999, rs.getInt("price")); + assertEquals(50, rs.getInt("quantity")); + assertTrue(rs.getBoolean("in_stock")); + } + } + + @Test + @DisplayName("Should throw DuplicateDocumentException when creating with existing key") + void testCreateDuplicateDocument() throws Exception { + // First create succeeds + ObjectNode objectNode = OBJECT_MAPPER.createObjectNode(); + objectNode.put("id", "dup-doc-200"); + objectNode.put("item", "First Item"); + Document document = new JSONDocument(objectNode); + Key key = new SingleValueKey("default", "dup-doc-200"); + + flatCollection.create(key, document); + + // Second create with same key should fail + ObjectNode objectNode2 = OBJECT_MAPPER.createObjectNode(); + objectNode2.put("id", "dup-doc-200"); + objectNode2.put("item", "Second Item"); + Document document2 = new JSONDocument(objectNode2); + + assertThrows(DuplicateDocumentException.class, () -> flatCollection.create(key, document2)); + } + + @Test + @DisplayName("Should create document with JSONB field") + void testCreateWithJsonbField() throws Exception { + ObjectNode objectNode = OBJECT_MAPPER.createObjectNode(); + objectNode.put("id", "jsonb-doc-300"); + objectNode.put("item", "Item with Props"); + ObjectNode propsNode = OBJECT_MAPPER.createObjectNode(); + propsNode.put("color", "blue"); + propsNode.put("size", "large"); + objectNode.set("props", propsNode); + Document document = new JSONDocument(objectNode); + Key key = new SingleValueKey("default", "jsonb-doc-300"); + + CreateResult result = flatCollection.create(key, document); + + assertTrue(result.isSucceed()); + + // Verify JSONB data + PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; + try (Connection conn = pgDatastore.getPostgresClient(); + PreparedStatement ps = + conn.prepareStatement( + String.format( + "SELECT props->>'color' as color FROM \"%s\" WHERE \"id\" = 'jsonb-doc-300'", + FLAT_COLLECTION_NAME)); + ResultSet rs = ps.executeQuery()) { + assertTrue(rs.next()); + assertEquals("blue", rs.getString("color")); + } + } + + @Test + @DisplayName("Should skip unknown fields when creating document") + void testCreateSkipsUnknownFields() throws Exception { + ObjectNode objectNode = OBJECT_MAPPER.createObjectNode(); + objectNode.put("id", "unknown-field-doc-400"); + objectNode.put("item", "Item"); + objectNode.put("unknown_column", "should be ignored"); + Document document = new JSONDocument(objectNode); + Key key = new SingleValueKey("default", "unknown-field-doc-400"); + + CreateResult result = flatCollection.create(key, document); - assertThrows(UnsupportedOperationException.class, () -> flatCollection.create(key, document)); + assertTrue(result.isSucceed()); } @Test @DisplayName("Should throw UnsupportedOperationException for createOrReplace") void testCreateOrReplace() { ObjectNode objectNode = OBJECT_MAPPER.createObjectNode(); - objectNode.put("_id", 200); + objectNode.put("id", "200"); objectNode.put("item", "NewMirror"); Document document = new JSONDocument(objectNode); Key key = new SingleValueKey("default", "200"); @@ -228,7 +319,7 @@ void testCreateOrReplace() { @DisplayName("Should throw UnsupportedOperationException for createOrReplaceAndReturn") void testCreateOrReplaceAndReturn() { ObjectNode objectNode = OBJECT_MAPPER.createObjectNode(); - objectNode.put("_id", 200); + objectNode.put("id", "200"); objectNode.put("item", "NewMirror"); Document document = new JSONDocument(objectNode); Key key = new SingleValueKey("default", "200"); diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java index 6dd89e1b..77a4c815 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java @@ -1,6 +1,17 @@ package org.hypertrace.core.documentstore.postgres; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; import java.io.IOException; +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.SQLException; +import java.sql.Statement; +import java.sql.Timestamp; +import java.sql.Types; +import java.time.Instant; +import java.util.ArrayList; +import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.Optional; @@ -15,12 +26,17 @@ import org.hypertrace.core.documentstore.Filter; import org.hypertrace.core.documentstore.Key; import org.hypertrace.core.documentstore.UpdateResult; +import org.hypertrace.core.documentstore.model.exception.DuplicateDocumentException; import org.hypertrace.core.documentstore.model.options.QueryOptions; import org.hypertrace.core.documentstore.model.options.UpdateOptions; import org.hypertrace.core.documentstore.model.subdoc.SubDocumentUpdate; +import org.hypertrace.core.documentstore.postgres.model.PostgresColumnMetadata; import org.hypertrace.core.documentstore.postgres.query.v1.PostgresQueryParser; +import org.hypertrace.core.documentstore.postgres.query.v1.parser.filter.nonjson.field.PostgresDataType; import org.hypertrace.core.documentstore.postgres.query.v1.transformer.FlatPostgresFieldTransformer; import org.hypertrace.core.documentstore.query.Query; +import org.postgresql.util.PSQLException; +import org.postgresql.util.PSQLState; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -34,6 +50,7 @@ public class FlatPostgresCollection extends PostgresCollection { private static final Logger LOGGER = LoggerFactory.getLogger(FlatPostgresCollection.class); + private static final ObjectMapper MAPPER = new ObjectMapper(); private static final String WRITE_NOT_SUPPORTED = "Write operations are not supported for flat collections yet!"; @@ -147,8 +164,173 @@ public void drop() { } @Override - public CreateResult create(Key key, Document document) throws IOException { - throw new UnsupportedOperationException(WRITE_NOT_SUPPORTED); + public CreateResult create(Key key, Document document) + throws IOException { + String tableName = tableIdentifier.getTableName(); + Map schema = schemaRegistry.getSchema(tableName); + + try { + JsonNode jsonNode = MAPPER.readTree(document.toJson()); + + List columns = new ArrayList<>(); + List values = new ArrayList<>(); + List types = new ArrayList<>(); + + Iterator> fields = jsonNode.fields(); + while (fields.hasNext()) { + Map.Entry field = fields.next(); + String fieldName = field.getKey(); + JsonNode fieldValue = field.getValue(); + + PostgresColumnMetadata columnMeta = schema.get(fieldName); + if (columnMeta == null) { + LOGGER.warn( + "Skipping field '{}' - not found in schema for table '{}'", fieldName, tableName); + continue; + } + + columns.add("\"" + fieldName + "\""); + values.add(extractValue(fieldValue, columnMeta.getPostgresType())); + types.add(columnMeta.getPostgresType()); + } + + if (columns.isEmpty()) { + throw new IOException("No valid columns found in document for table: " + tableName); + } + + String sql = buildInsertSql(columns); + LOGGER.debug("Insert SQL: {}", sql); + + try (Connection conn = client.getPooledConnection(); + PreparedStatement ps = conn.prepareStatement(sql, Statement.RETURN_GENERATED_KEYS)) { + + for (int i = 0; i < values.size(); i++) { + setParameter(ps, i + 1, values.get(i), types.get(i)); + } + + int result = ps.executeUpdate(); + LOGGER.debug("Create result: {}", result); + return new CreateResult(result > 0); + } + } catch (PSQLException e) { + if (PSQLState.UNIQUE_VIOLATION.getState().equals(e.getSQLState())) { + throw new DuplicateDocumentException(); + } + LOGGER.error("SQLException creating document. key: {} content: {}", key, document, e); + throw new IOException(e); + } catch (SQLException e) { + LOGGER.error("SQLException creating document. key: {} content: {}", key, document, e); + throw new IOException(e); + } + } + + private String buildInsertSql(List columns) { + String columnList = String.join(", ", columns); + String placeholders = String.join(", ", columns.stream().map(c -> "?").toArray(String[]::new)); + return String.format( + "INSERT INTO %s (%s) VALUES (%s)", tableIdentifier, columnList, placeholders); + } + + private Object extractValue(JsonNode node, PostgresDataType type) { + if (node == null || node.isNull()) { + return null; + } + + switch (type) { + case INTEGER: + return node.isNumber() ? node.intValue() : Integer.parseInt(node.asText()); + case BIGINT: + return node.isNumber() ? node.longValue() : Long.parseLong(node.asText()); + case REAL: + return node.isNumber() ? node.floatValue() : Float.parseFloat(node.asText()); + case DOUBLE_PRECISION: + return node.isNumber() ? node.doubleValue() : Double.parseDouble(node.asText()); + case BOOLEAN: + return node.isBoolean() ? node.booleanValue() : Boolean.parseBoolean(node.asText()); + case TEXT: + return node.asText(); + case TIMESTAMPTZ: + if (node.isTextual()) { + return Timestamp.from(Instant.parse(node.asText())); + } else if (node.isNumber()) { + return new Timestamp(node.longValue()); + } + return null; + case DATE: + if (node.isTextual()) { + return java.sql.Date.valueOf(node.asText()); + } + return null; + case JSONB: + return node.toString(); + default: + return node.asText(); + } + } + + private void setParameter(PreparedStatement ps, int index, Object value, PostgresDataType type) + throws SQLException { + if (value == null) { + ps.setNull(index, getSqlType(type)); + return; + } + + switch (type) { + case INTEGER: + ps.setInt(index, (Integer) value); + break; + case BIGINT: + ps.setLong(index, (Long) value); + break; + case REAL: + ps.setFloat(index, (Float) value); + break; + case DOUBLE_PRECISION: + ps.setDouble(index, (Double) value); + break; + case BOOLEAN: + ps.setBoolean(index, (Boolean) value); + break; + case TEXT: + ps.setString(index, (String) value); + break; + case TIMESTAMPTZ: + ps.setTimestamp(index, (Timestamp) value); + break; + case DATE: + ps.setDate(index, (java.sql.Date) value); + break; + case JSONB: + ps.setObject(index, value, Types.OTHER); + break; + default: + ps.setString(index, value.toString()); + } + } + + private int getSqlType(PostgresDataType type) { + switch (type) { + case INTEGER: + return Types.INTEGER; + case BIGINT: + return Types.BIGINT; + case REAL: + return Types.REAL; + case DOUBLE_PRECISION: + return Types.DOUBLE; + case BOOLEAN: + return Types.BOOLEAN; + case TEXT: + return Types.VARCHAR; + case TIMESTAMPTZ: + return Types.TIMESTAMP_WITH_TIMEZONE; + case DATE: + return Types.DATE; + case JSONB: + return Types.OTHER; + default: + return Types.VARCHAR; + } } @Override From 31f16e27a1151148370e989a0f6a900c5166df01 Mon Sep 17 00:00:00 2001 From: Prashant Pandey Date: Sun, 4 Jan 2026 19:48:51 +0530 Subject: [PATCH 16/26] Refactor --- ...yLoadedSchemaRegistryIntegrationTest.java} | 2 +- .../documentstore/commons/ColumnMetadata.java | 10 +++++++ .../documentstore/commons/SchemaRegistry.java | 28 +++++++++++++++++++ .../postgres/FlatPostgresCollection.java | 4 +-- .../postgres/PostgresDatastore.java | 4 +-- ... PostgresLazyilyLoadedSchemaRegistry.java} | 21 ++++++++++++-- .../postgres/PostgresMetadataFetcher.java | 6 +--- .../model/PostgresColumnMetadata.java | 1 - ...tgresLazyilyLoadedSchemaRegistryTest.java} | 11 ++------ 9 files changed, 65 insertions(+), 22 deletions(-) rename document-store/src/integrationTest/java/org/hypertrace/core/documentstore/postgres/{PostgresSchemaRegistryIntegrationTest.java => PostgresLazyilyLoadedSchemaRegistryIntegrationTest.java} (99%) rename document-store/src/main/java/org/hypertrace/core/documentstore/postgres/{PostgresSchemaRegistry.java => PostgresLazyilyLoadedSchemaRegistry.java} (84%) rename document-store/src/test/java/org/hypertrace/core/documentstore/postgres/{PostgresSchemaRegistryTest.java => PostgresLazyilyLoadedSchemaRegistryTest.java} (97%) diff --git a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistryIntegrationTest.java b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/postgres/PostgresLazyilyLoadedSchemaRegistryIntegrationTest.java similarity index 99% rename from document-store/src/integrationTest/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistryIntegrationTest.java rename to document-store/src/integrationTest/java/org/hypertrace/core/documentstore/postgres/PostgresLazyilyLoadedSchemaRegistryIntegrationTest.java index 2b3b30d4..53cd9f28 100644 --- a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistryIntegrationTest.java +++ b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/postgres/PostgresLazyilyLoadedSchemaRegistryIntegrationTest.java @@ -29,7 +29,7 @@ import org.testcontainers.utility.DockerImageName; @Testcontainers -class PostgresSchemaRegistryIntegrationTest { +class PostgresLazyilyLoadedSchemaRegistryIntegrationTest { private static final String TABLE_NAME = "myTestFlat"; diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/commons/ColumnMetadata.java b/document-store/src/main/java/org/hypertrace/core/documentstore/commons/ColumnMetadata.java index 6fa89dc3..475e55be 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/commons/ColumnMetadata.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/commons/ColumnMetadata.java @@ -3,9 +3,19 @@ import org.hypertrace.core.documentstore.expression.impl.DataType; public interface ColumnMetadata { + + /** + * @return the col name + */ String getName(); + /** + * @return the col's canonical type, as defined here: {@link DataType} + */ DataType getCanonicalType(); + /** + * @return whether this column can be set to NULL + */ boolean isNullable(); } diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/commons/SchemaRegistry.java b/document-store/src/main/java/org/hypertrace/core/documentstore/commons/SchemaRegistry.java index 68782184..fd3fa792 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/commons/SchemaRegistry.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/commons/SchemaRegistry.java @@ -3,11 +3,39 @@ import java.util.Map; import java.util.Optional; +/** + * SchemaRegistry is an interface for a registry of schemas. This interface does not places any + * restrictions on how schemas are loaded. They can be loaded at bootstrap and cached, or loaded + * lazily, or via any other method. + * + * @param the type of metadata for a particular column in the registry + */ public interface SchemaRegistry { + /** + * Returns the schema for a particular table. If the schema is not available for that table, it + * returns null (note that some implementations may fetch the schema if this happens. The + * interface does not make it mandatory). + * + * @param tableName the table for which schema has to be fetched + * @return a map of column name to their metadata + */ Map getSchema(String tableName); + /** + * Invalidates the current schema of the table that the schema registry is holding + * + * @param tableName the table name + */ void invalidate(String tableName); + /** + * Returns the metadata of a col from the registry. If the metadata is not found, an + * implementation might fetch it from the source synchronously. + * + * @param tableName the table name + * @param colName the col name + * @return optional of the col metadata. + */ Optional getColumnOrRefresh(String tableName, String colName); } diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java index 6dd89e1b..747b63bd 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java @@ -37,12 +37,12 @@ public class FlatPostgresCollection extends PostgresCollection { private static final String WRITE_NOT_SUPPORTED = "Write operations are not supported for flat collections yet!"; - private final PostgresSchemaRegistry schemaRegistry; + private final PostgresLazyilyLoadedSchemaRegistry schemaRegistry; FlatPostgresCollection( final PostgresClient client, final String collectionName, - final PostgresSchemaRegistry schemaRegistry) { + final PostgresLazyilyLoadedSchemaRegistry schemaRegistry) { super(client, collectionName); this.schemaRegistry = schemaRegistry; } 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 4e9fb3e2..e2d6ac58 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 @@ -61,7 +61,7 @@ public PostgresDatastore(@NonNull final DatastoreConfig datastoreConfig) { database = connectionConfig.database(); docStoreMetricProvider = new PostgresDocStoreMetricProvider(this, postgresConnectionConfig); schemaRegistry = - new PostgresSchemaRegistry( + new PostgresLazyilyLoadedSchemaRegistry( new PostgresMetadataFetcher(client), postgresConnectionConfig.schemaCacheExpiry(), postgresConnectionConfig.schemaRefreshCooldown()); @@ -172,7 +172,7 @@ public Collection getCollectionForType(String collectionName, DocumentType docum case FLAT: { return new FlatPostgresCollection( - client, collectionName, (PostgresSchemaRegistry) schemaRegistry); + client, collectionName, (PostgresLazyilyLoadedSchemaRegistry) schemaRegistry); } case NESTED: return getCollection(collectionName); diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistry.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresLazyilyLoadedSchemaRegistry.java similarity index 84% rename from document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistry.java rename to document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresLazyilyLoadedSchemaRegistry.java index e8fc9b19..38ac8dd5 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistry.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresLazyilyLoadedSchemaRegistry.java @@ -12,6 +12,8 @@ import java.util.concurrent.TimeUnit; import org.hypertrace.core.documentstore.commons.SchemaRegistry; import org.hypertrace.core.documentstore.postgres.model.PostgresColumnMetadata; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** * A lazily-loaded, cached schema registry for PostgreSQL tables. @@ -49,10 +51,17 @@ * @see PostgresMetadataFetcher * @see PostgresColumnMetadata */ -public class PostgresSchemaRegistry implements SchemaRegistry { +public class PostgresLazyilyLoadedSchemaRegistry implements SchemaRegistry { + private static final Logger LOGGER = + LoggerFactory.getLogger(PostgresLazyilyLoadedSchemaRegistry.class); + + // The cache registry - Key: Table name, value: Map of column name to column metadata private final LoadingCache> cache; + // This tracks when was the last time a cache was refreshed for a table. This helps track the + // cooldown period on a per-table level private final ConcurrentHashMap lastRefreshTimes; + // How long to wait for a tables' schema refresh before it is refreshed again private final Duration refreshCooldown; /** @@ -62,8 +71,9 @@ public class PostgresSchemaRegistry implements SchemaRegistry(); this.cache = @@ -73,8 +83,11 @@ public PostgresSchemaRegistry( new CacheLoader<>() { @Override public Map load(String tableName) { + LOGGER.info("Loading schema for table: {}", tableName); + Map updatedSchema = fetcher.fetch(tableName); lastRefreshTimes.put(tableName, Instant.now()); - return fetcher.fetch(tableName); + LOGGER.info("Successfully loading schema for table: {}", tableName); + return updatedSchema; } }); } @@ -91,6 +104,7 @@ public Map getSchema(String tableName) { try { return cache.get(tableName); } catch (ExecutionException e) { + LOGGER.error("Could not fetch the schema for table from the cache: {}", tableName, e); throw new RuntimeException("Failed to fetch schema for " + tableName, e.getCause()); } } @@ -104,6 +118,7 @@ public Map getSchema(String tableName) { */ @Override public void invalidate(String tableName) { + LOGGER.info("Invalidating schema for table: {}", tableName); cache.invalidate(tableName); } diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresMetadataFetcher.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresMetadataFetcher.java index 2d2be0ec..24b10e07 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresMetadataFetcher.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresMetadataFetcher.java @@ -40,11 +40,7 @@ public Map fetch(String tableName) { metadataMap.put( columnName, new PostgresColumnMetadata( - columnName, - mapToCanonicalType(udtName), - mapToPostgresType(udtName), - udtName, - isNullable)); + columnName, mapToCanonicalType(udtName), mapToPostgresType(udtName), isNullable)); } } return metadataMap; diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/model/PostgresColumnMetadata.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/model/PostgresColumnMetadata.java index 508b3dc6..3a2d4540 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/model/PostgresColumnMetadata.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/model/PostgresColumnMetadata.java @@ -14,7 +14,6 @@ public class PostgresColumnMetadata implements ColumnMetadata { private final String colName; private final DataType canonicalType; @Getter private final PostgresDataType postgresType; - private final String pgType; private final boolean nullable; @Override diff --git a/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistryTest.java b/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresLazyilyLoadedSchemaRegistryTest.java similarity index 97% rename from document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistryTest.java rename to document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresLazyilyLoadedSchemaRegistryTest.java index 7fa22d5f..4257f8fe 100644 --- a/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresSchemaRegistryTest.java +++ b/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresLazyilyLoadedSchemaRegistryTest.java @@ -22,7 +22,7 @@ import org.mockito.junit.jupiter.MockitoExtension; @ExtendWith(MockitoExtension.class) -class PostgresSchemaRegistryTest { +class PostgresLazyilyLoadedSchemaRegistryTest { private static final String TEST_TABLE = "test_table"; private static final String COL_ID = "id"; @@ -33,11 +33,11 @@ class PostgresSchemaRegistryTest { @Mock private PostgresMetadataFetcher fetcher; - private PostgresSchemaRegistry registry; + private PostgresLazyilyLoadedSchemaRegistry registry; @BeforeEach void setUp() { - registry = new PostgresSchemaRegistry(fetcher, CACHE_EXPIRY, REFRESH_COOLDOWN); + registry = new PostgresLazyilyLoadedSchemaRegistry(fetcher, CACHE_EXPIRY, REFRESH_COOLDOWN); } @Test @@ -79,7 +79,6 @@ void getSchemaLoadsEachTableIndependently() { .colName("other_col") .canonicalType(DataType.BOOLEAN) .postgresType(PostgresDataType.BOOLEAN) - .pgType("bool") .nullable(false) .build()); @@ -164,7 +163,6 @@ void getColumnOrRefreshRefreshesSchemaIfColumnMissingAndCooldownExpired() throws .colName("new_col") .canonicalType(DataType.STRING) .postgresType(PostgresDataType.TEXT) - .pgType("text") .nullable(true) .build()); @@ -280,7 +278,6 @@ private Map createTestSchema() { .colName(COL_ID) .canonicalType(DataType.INTEGER) .postgresType(PostgresDataType.INTEGER) - .pgType("int4") .nullable(false) .build()); schema.put( @@ -289,7 +286,6 @@ private Map createTestSchema() { .colName(COL_NAME) .canonicalType(DataType.STRING) .postgresType(PostgresDataType.TEXT) - .pgType("text") .nullable(true) .build()); schema.put( @@ -298,7 +294,6 @@ private Map createTestSchema() { .colName(COL_PRICE) .canonicalType(DataType.DOUBLE) .postgresType(PostgresDataType.DOUBLE_PRECISION) - .pgType("float8") .nullable(true) .build()); return schema; From bf5ca5ce900ab8beae72906544883d9b8ef4762f Mon Sep 17 00:00:00 2001 From: Prashant Pandey Date: Mon, 12 Jan 2026 12:35:04 +0530 Subject: [PATCH 17/26] WIP --- .../FlatCollectionWriteTest.java | 133 +++++- .../postgres/FlatPostgresCollection.java | 413 +++++++++++++----- 2 files changed, 422 insertions(+), 124 deletions(-) diff --git a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java index 65d26904..68051082 100644 --- a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java +++ b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java @@ -2,6 +2,7 @@ import static org.hypertrace.core.documentstore.utils.Utils.readFileFromResource; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -288,45 +289,145 @@ void testCreateWithJsonbField() throws Exception { } @Test - @DisplayName("Should skip unknown fields when creating document") + @DisplayName("Should skip unknown fields and insert known fields") void testCreateSkipsUnknownFields() throws Exception { ObjectNode objectNode = OBJECT_MAPPER.createObjectNode(); objectNode.put("id", "unknown-field-doc-400"); objectNode.put("item", "Item"); - objectNode.put("unknown_column", "should be ignored"); + objectNode.put("unknown_column", "should be skipped"); Document document = new JSONDocument(objectNode); Key key = new SingleValueKey("default", "unknown-field-doc-400"); CreateResult result = flatCollection.create(key, document); assertTrue(result.isSucceed()); + + // Verify only known columns were inserted + PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; + try (Connection conn = pgDatastore.getPostgresClient(); + PreparedStatement ps = + conn.prepareStatement( + String.format( + "SELECT * FROM \"%s\" WHERE \"id\" = 'unknown-field-doc-400'", + FLAT_COLLECTION_NAME)); + ResultSet rs = ps.executeQuery()) { + assertTrue(rs.next()); + assertEquals("Item", rs.getString("item")); + } } @Test - @DisplayName("Should throw UnsupportedOperationException for createOrReplace") - void testCreateOrReplace() { + @DisplayName("Should create a new document when key does not exist") + void testCreateOrReplaceNewDocument() throws Exception { ObjectNode objectNode = OBJECT_MAPPER.createObjectNode(); - objectNode.put("id", "200"); + objectNode.put("id", "createorreplace-new-500"); objectNode.put("item", "NewMirror"); + objectNode.put("price", 150); Document document = new JSONDocument(objectNode); - Key key = new SingleValueKey("default", "200"); + Key key = new SingleValueKey("default", "createorreplace-new-500"); - assertThrows( - UnsupportedOperationException.class, () -> flatCollection.createOrReplace(key, document)); + boolean result = flatCollection.createOrReplace(key, document); + + assertTrue(result); + + // Verify the data was inserted + PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; + try (Connection conn = pgDatastore.getPostgresClient(); + PreparedStatement ps = + conn.prepareStatement( + String.format( + "SELECT * FROM \"%s\" WHERE \"id\" = 'createorreplace-new-500'", + FLAT_COLLECTION_NAME)); + ResultSet rs = ps.executeQuery()) { + assertTrue(rs.next()); + assertEquals("NewMirror", rs.getString("item")); + assertEquals(150, rs.getInt("price")); + } } @Test - @DisplayName("Should throw UnsupportedOperationException for createOrReplaceAndReturn") - void testCreateOrReplaceAndReturn() { + @DisplayName("Should replace an existing document when key exists") + void testCreateOrReplaceExistingDocument() throws Exception { + // First create a document + ObjectNode objectNode1 = OBJECT_MAPPER.createObjectNode(); + objectNode1.put("id", "createorreplace-existing-600"); + objectNode1.put("item", "OriginalItem"); + objectNode1.put("price", 100); + Document document1 = new JSONDocument(objectNode1); + Key key = new SingleValueKey("default", "createorreplace-existing-600"); + flatCollection.create(key, document1); + + // Now replace it + ObjectNode objectNode2 = OBJECT_MAPPER.createObjectNode(); + objectNode2.put("id", "createorreplace-existing-600"); + objectNode2.put("item", "ReplacedItem"); + objectNode2.put("price", 200); + Document document2 = new JSONDocument(objectNode2); + + boolean result = flatCollection.createOrReplace(key, document2); + + assertTrue(result); + + // Verify the data was replaced + PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; + try (Connection conn = pgDatastore.getPostgresClient(); + PreparedStatement ps = + conn.prepareStatement( + String.format( + "SELECT * FROM \"%s\" WHERE \"id\" = 'createorreplace-existing-600'", + FLAT_COLLECTION_NAME)); + ResultSet rs = ps.executeQuery()) { + assertTrue(rs.next()); + assertEquals("ReplacedItem", rs.getString("item")); + assertEquals(200, rs.getInt("price")); + } + } + + @Test + @DisplayName("Should create and return a new document when key does not exist") + void testCreateOrReplaceAndReturnNewDocument() throws Exception { ObjectNode objectNode = OBJECT_MAPPER.createObjectNode(); - objectNode.put("id", "200"); - objectNode.put("item", "NewMirror"); + objectNode.put("id", "createorreplace-return-700"); + objectNode.put("item", "ReturnedItem"); + objectNode.put("price", 250); Document document = new JSONDocument(objectNode); - Key key = new SingleValueKey("default", "200"); + Key key = new SingleValueKey("default", "createorreplace-return-700"); - assertThrows( - UnsupportedOperationException.class, - () -> flatCollection.createOrReplaceAndReturn(key, document)); + Document result = flatCollection.createOrReplaceAndReturn(key, document); + + assertNotNull(result); + JsonNode resultNode = OBJECT_MAPPER.readTree(result.toJson()); + assertEquals("createorreplace-return-700", resultNode.get("id").asText()); + assertEquals("ReturnedItem", resultNode.get("item").asText()); + assertEquals(250, resultNode.get("price").asInt()); + } + + @Test + @DisplayName("Should replace and return an existing document when key exists") + void testCreateOrReplaceAndReturnExistingDocument() throws Exception { + // First create a document + ObjectNode objectNode1 = OBJECT_MAPPER.createObjectNode(); + objectNode1.put("id", "createorreplace-return-existing-800"); + objectNode1.put("item", "OriginalReturnItem"); + objectNode1.put("price", 300); + Document document1 = new JSONDocument(objectNode1); + Key key = new SingleValueKey("default", "createorreplace-return-existing-800"); + flatCollection.create(key, document1); + + // Now replace it + ObjectNode objectNode2 = OBJECT_MAPPER.createObjectNode(); + objectNode2.put("id", "createorreplace-return-existing-800"); + objectNode2.put("item", "ReplacedReturnItem"); + objectNode2.put("price", 400); + Document document2 = new JSONDocument(objectNode2); + + Document result = flatCollection.createOrReplaceAndReturn(key, document2); + + assertNotNull(result); + JsonNode resultNode = OBJECT_MAPPER.readTree(result.toJson()); + assertEquals("createorreplace-return-existing-800", resultNode.get("id").asText()); + assertEquals("ReplacedReturnItem", resultNode.get("item").asText()); + assertEquals(400, resultNode.get("price").asInt()); } } diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java index d2bc6b75..6dedceef 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java @@ -4,9 +4,10 @@ import com.fasterxml.jackson.databind.ObjectMapper; import java.io.IOException; import java.sql.Connection; +import java.sql.Date; import java.sql.PreparedStatement; +import java.sql.ResultSet; import java.sql.SQLException; -import java.sql.Statement; import java.sql.Timestamp; import java.sql.Types; import java.time.Instant; @@ -14,8 +15,10 @@ import java.util.Iterator; import java.util.List; import java.util.Map; +import java.util.Map.Entry; import java.util.Optional; import java.util.Set; +import java.util.stream.Collectors; import org.hypertrace.core.documentstore.BulkArrayValueUpdateRequest; import org.hypertrace.core.documentstore.BulkDeleteResult; import org.hypertrace.core.documentstore.BulkUpdateRequest; @@ -24,6 +27,7 @@ import org.hypertrace.core.documentstore.CreateResult; import org.hypertrace.core.documentstore.Document; import org.hypertrace.core.documentstore.Filter; +import org.hypertrace.core.documentstore.JSONDocument; import org.hypertrace.core.documentstore.Key; import org.hypertrace.core.documentstore.UpdateResult; import org.hypertrace.core.documentstore.model.exception.DuplicateDocumentException; @@ -164,66 +168,324 @@ public void drop() { } @Override - public CreateResult create(Key key, Document document) + public CreateResult create(Key key, Document document) throws IOException { + return createWithRetry(key, document, false); + } + + @Override + public boolean createOrReplace(Key key, Document document) throws IOException { + return createOrReplaceWithRetry(key, document, false); + } + + @Override + public Document createOrReplaceAndReturn(Key key, Document document) throws IOException { + return createOrReplaceAndReturnWithRetry(key, document, false); + } + + @Override + public BulkUpdateResult bulkUpdate(List bulkUpdateRequests) { + throw new UnsupportedOperationException(WRITE_NOT_SUPPORTED); + } + + @Override + public UpdateResult update(Key key, Document document, Filter condition) throws IOException { + throw new UnsupportedOperationException(WRITE_NOT_SUPPORTED); + } + + @Override + public Optional update( + org.hypertrace.core.documentstore.query.Query query, + java.util.Collection updates, + UpdateOptions updateOptions) throws IOException { - String tableName = tableIdentifier.getTableName(); - Map schema = schemaRegistry.getSchema(tableName); + throw new UnsupportedOperationException(WRITE_NOT_SUPPORTED); + } - try { - JsonNode jsonNode = MAPPER.readTree(document.toJson()); - - List columns = new ArrayList<>(); - List values = new ArrayList<>(); - List types = new ArrayList<>(); - - Iterator> fields = jsonNode.fields(); - while (fields.hasNext()) { - Map.Entry field = fields.next(); - String fieldName = field.getKey(); - JsonNode fieldValue = field.getValue(); - - PostgresColumnMetadata columnMeta = schema.get(fieldName); - if (columnMeta == null) { - LOGGER.warn( - "Skipping field '{}' - not found in schema for table '{}'", fieldName, tableName); - continue; - } + @Override + public CloseableIterator bulkUpdate( + org.hypertrace.core.documentstore.query.Query query, + java.util.Collection updates, + UpdateOptions updateOptions) + throws IOException { + throw new UnsupportedOperationException(WRITE_NOT_SUPPORTED); + } - columns.add("\"" + fieldName + "\""); - values.add(extractValue(fieldValue, columnMeta.getPostgresType())); - types.add(columnMeta.getPostgresType()); - } + private CreateResult createWithRetry(Key key, Document document, boolean isRetry) + throws IOException { + String tableName = tableIdentifier.getTableName(); - if (columns.isEmpty()) { - throw new IOException("No valid columns found in document for table: " + tableName); + try { + TypedDocument parsed = parseDocument(document, tableName); + // if there are no valid columns in the document + if (parsed.isEmpty()) { + LOGGER.warn("No valid columns found in the document for table: {}", tableName); + return new CreateResult(false); } - String sql = buildInsertSql(columns); + String sql = buildInsertSql(parsed.getColumns()); LOGGER.debug("Insert SQL: {}", sql); - try (Connection conn = client.getPooledConnection(); - PreparedStatement ps = conn.prepareStatement(sql, Statement.RETURN_GENERATED_KEYS)) { + int result = executeUpdate(sql, parsed); + LOGGER.debug("Create result: {}", result); + return new CreateResult(result > 0); - for (int i = 0; i < values.size(); i++) { - setParameter(ps, i + 1, values.get(i), types.get(i)); - } - - int result = ps.executeUpdate(); - LOGGER.debug("Create result: {}", result); - return new CreateResult(result > 0); - } } catch (PSQLException e) { if (PSQLState.UNIQUE_VIOLATION.getState().equals(e.getSQLState())) { throw new DuplicateDocumentException(); } + return handlePSQLExceptionForCreate(e, key, document, tableName, isRetry); + } catch (SQLException e) { LOGGER.error("SQLException creating document. key: {} content: {}", key, document, e); throw new IOException(e); + } + } + + private boolean createOrReplaceWithRetry(Key key, Document document, boolean isRetry) + throws IOException { + String tableName = tableIdentifier.getTableName(); + + try { + TypedDocument parsed = parseDocument(document, tableName); + if (parsed.isEmpty()) { + LOGGER.warn("No valid columns found in the document for table: {}", tableName); + return false; + } + + String sql = buildUpsertSql(parsed.getColumns()); + LOGGER.debug("Upsert SQL: {}", sql); + + int result = executeUpdate(sql, parsed); + LOGGER.debug("CreateOrReplace result: {}", result); + return result > 0; + + } catch (PSQLException e) { + return handlePSQLExceptionForCreateOrReplace(e, key, document, tableName, isRetry); } catch (SQLException e) { - LOGGER.error("SQLException creating document. key: {} content: {}", key, document, e); + LOGGER.error("SQLException in createOrReplace. key: {} content: {}", key, document, e); throw new IOException(e); } } + private Document createOrReplaceAndReturnWithRetry(Key key, Document document, boolean isRetry) + throws IOException { + String tableName = tableIdentifier.getTableName(); + + try { + TypedDocument parsed = parseDocument(document, tableName); + if (parsed.isEmpty()) { + LOGGER.warn("No valid columns found in the document for table: {}", tableName); + return null; + } + + String sql = buildUpsertSqlWithReturning(parsed.getColumns()); + LOGGER.debug("Upsert with RETURNING SQL: {}", sql); + + return executeQueryAndReturn(sql, parsed); + + } catch (PSQLException e) { + return handlePSQLExceptionForCreateOrReplaceAndReturn(e, key, document, tableName, isRetry); + } catch (SQLException e) { + LOGGER.error( + "SQLException in createOrReplaceAndReturn. key: {} content: {}", key, document, e); + throw new IOException(e); + } + } + + private TypedDocument parseDocument(Document document, String tableName) throws IOException { + JsonNode jsonNode = MAPPER.readTree(document.toJson()); + List entries = new ArrayList<>(); + + Iterator> fields = jsonNode.fields(); + while (fields.hasNext()) { + Entry field = fields.next(); + String fieldName = field.getKey(); + JsonNode fieldValue = field.getValue(); + + Optional columnMetadata = + schemaRegistry.getColumnOrRefresh(tableName, fieldName); + + if (columnMetadata.isEmpty()) { + LOGGER.warn("Could not find column metadata for column: {}, skipping it", fieldName); + continue; + } + + PostgresDataType type = columnMetadata.get().getPostgresType(); + entries.add(new ColumnEntry("\"" + fieldName + "\"", extractValue(fieldValue, type), type)); + } + + return new TypedDocument(entries); + } + + private int executeUpdate(String sql, TypedDocument parsed) throws SQLException { + try (Connection conn = client.getPooledConnection(); + PreparedStatement ps = conn.prepareStatement(sql)) { + int index = 1; + for (ColumnEntry entry : parsed.entries) { + setParameter(ps, index++, entry.value, entry.type); + } + return ps.executeUpdate(); + } + } + + private Document executeQueryAndReturn(String sql, TypedDocument parsed) + throws SQLException, IOException { + try (Connection conn = client.getPooledConnection(); + PreparedStatement ps = conn.prepareStatement(sql)) { + int index = 1; + for (ColumnEntry entry : parsed.entries) { + setParameter(ps, index++, entry.value, entry.type); + } + try (ResultSet rs = ps.executeQuery()) { + if (rs.next()) { + return resultSetToDocument(rs); + } + return null; + } + } + } + + private CreateResult handlePSQLExceptionForCreate( + PSQLException e, Key key, Document document, String tableName, boolean isRetry) + throws IOException { + if (!isRetry && shouldRefreshSchemaAndRetry(e.getSQLState())) { + LOGGER.info( + "Schema mismatch detected (SQLState: {}), refreshing schema and retrying. key: {}", + e.getSQLState(), + key); + schemaRegistry.invalidate(tableName); + return createWithRetry(key, document, true); + } + LOGGER.error("SQLException creating document. key: {} content: {}", key, document, e); + throw new IOException(e); + } + + private boolean handlePSQLExceptionForCreateOrReplace( + PSQLException e, Key key, Document document, String tableName, boolean isRetry) + throws IOException { + if (!isRetry && shouldRefreshSchemaAndRetry(e.getSQLState())) { + LOGGER.info( + "Schema mismatch detected (SQLState: {}), refreshing schema and retrying. key: {}", + e.getSQLState(), + key); + schemaRegistry.invalidate(tableName); + return createOrReplaceWithRetry(key, document, true); + } + LOGGER.error("SQLException in createOrReplace. key: {} content: {}", key, document, e); + throw new IOException(e); + } + + private Document handlePSQLExceptionForCreateOrReplaceAndReturn( + PSQLException e, Key key, Document document, String tableName, boolean isRetry) + throws IOException { + if (!isRetry && shouldRefreshSchemaAndRetry(e.getSQLState())) { + LOGGER.info( + "Schema mismatch detected (SQLState: {}), refreshing schema and retrying. key: {}", + e.getSQLState(), + key); + schemaRegistry.invalidate(tableName); + return createOrReplaceAndReturnWithRetry(key, document, true); + } + LOGGER.error("SQLException in createOrReplaceAndReturn. key: {} content: {}", key, document, e); + throw new IOException(e); + } + + private boolean shouldRefreshSchemaAndRetry(String sqlState) { + return PSQLState.UNDEFINED_COLUMN.getState().equals(sqlState) + || PSQLState.DATATYPE_MISMATCH.getState().equals(sqlState); + } + + private static class ColumnEntry { + final String column; + final Object value; + final PostgresDataType type; + + ColumnEntry(String column, Object value, PostgresDataType type) { + this.column = column; + this.value = value; + this.type = type; + } + } + + private static class TypedDocument { + final List entries; + + TypedDocument(List entries) { + this.entries = entries; + } + + boolean isEmpty() { + return entries.isEmpty(); + } + + List getColumns() { + return entries.stream().map(e -> e.column).collect(Collectors.toList()); + } + } + + private String buildUpsertSql(List columns) { + String columnList = String.join(", ", columns); + String placeholders = String.join(", ", columns.stream().map(c -> "?").toArray(String[]::new)); + String updateSet = + columns.stream() + .filter(col -> !"\"id\"".equals(col)) + .map(col -> col + " = EXCLUDED." + col) + .collect(Collectors.joining(", ")); + + if (updateSet.isEmpty()) { + return String.format( + "INSERT INTO %s (%s) VALUES (%s) ON CONFLICT (\"id\") DO NOTHING", + tableIdentifier, columnList, placeholders); + } + + return String.format( + "INSERT INTO %s (%s) VALUES (%s) ON CONFLICT (\"id\") DO UPDATE SET %s", + tableIdentifier, columnList, placeholders, updateSet); + } + + private String buildUpsertSqlWithReturning(List columns) { + return buildUpsertSql(columns) + " RETURNING *"; + } + + private Document resultSetToDocument(ResultSet rs) throws SQLException, IOException { + int columnCount = rs.getMetaData().getColumnCount(); + com.fasterxml.jackson.databind.node.ObjectNode objectNode = MAPPER.createObjectNode(); + + for (int i = 1; i <= columnCount; i++) { + String columnName = rs.getMetaData().getColumnName(i); + Object value = rs.getObject(i); + + if (value == null) { + objectNode.putNull(columnName); + } else if (value instanceof Integer) { + objectNode.put(columnName, (Integer) value); + } else if (value instanceof Long) { + objectNode.put(columnName, (Long) value); + } else if (value instanceof Double) { + objectNode.put(columnName, (Double) value); + } else if (value instanceof Float) { + objectNode.put(columnName, (Float) value); + } else if (value instanceof Boolean) { + objectNode.put(columnName, (Boolean) value); + } else if (value instanceof Timestamp) { + objectNode.put(columnName, ((Timestamp) value).toInstant().toString()); + } else if (value instanceof java.sql.Date) { + objectNode.put(columnName, value.toString()); + } else if (value instanceof org.postgresql.util.PGobject) { + // Handle JSONB + String jsonValue = ((org.postgresql.util.PGobject) value).getValue(); + if (jsonValue != null) { + objectNode.set(columnName, MAPPER.readTree(jsonValue)); + } else { + objectNode.putNull(columnName); + } + } else { + objectNode.put(columnName, value.toString()); + } + } + + return new JSONDocument(objectNode); + } + private String buildInsertSql(List columns) { String columnList = String.join(", ", columns); String placeholders = String.join(", ", columns.stream().map(c -> "?").toArray(String[]::new)); @@ -247,8 +509,6 @@ private Object extractValue(JsonNode node, PostgresDataType type) { return node.isNumber() ? node.doubleValue() : Double.parseDouble(node.asText()); case BOOLEAN: return node.isBoolean() ? node.booleanValue() : Boolean.parseBoolean(node.asText()); - case TEXT: - return node.asText(); case TIMESTAMPTZ: if (node.isTextual()) { return Timestamp.from(Instant.parse(node.asText())); @@ -258,7 +518,7 @@ private Object extractValue(JsonNode node, PostgresDataType type) { return null; case DATE: if (node.isTextual()) { - return java.sql.Date.valueOf(node.asText()); + return Date.valueOf(node.asText()); } return null; case JSONB: @@ -271,7 +531,7 @@ private Object extractValue(JsonNode node, PostgresDataType type) { private void setParameter(PreparedStatement ps, int index, Object value, PostgresDataType type) throws SQLException { if (value == null) { - ps.setNull(index, getSqlType(type)); + ps.setObject(index, null); return; } @@ -307,67 +567,4 @@ private void setParameter(PreparedStatement ps, int index, Object value, Postgre ps.setString(index, value.toString()); } } - - private int getSqlType(PostgresDataType type) { - switch (type) { - case INTEGER: - return Types.INTEGER; - case BIGINT: - return Types.BIGINT; - case REAL: - return Types.REAL; - case DOUBLE_PRECISION: - return Types.DOUBLE; - case BOOLEAN: - return Types.BOOLEAN; - case TEXT: - return Types.VARCHAR; - case TIMESTAMPTZ: - return Types.TIMESTAMP_WITH_TIMEZONE; - case DATE: - return Types.DATE; - case JSONB: - return Types.OTHER; - default: - return Types.VARCHAR; - } - } - - @Override - public boolean createOrReplace(Key key, Document document) throws IOException { - throw new UnsupportedOperationException(WRITE_NOT_SUPPORTED); - } - - @Override - public Document createOrReplaceAndReturn(Key key, Document document) throws IOException { - throw new UnsupportedOperationException(WRITE_NOT_SUPPORTED); - } - - @Override - public BulkUpdateResult bulkUpdate(List bulkUpdateRequests) { - throw new UnsupportedOperationException(WRITE_NOT_SUPPORTED); - } - - @Override - public UpdateResult update(Key key, Document document, Filter condition) throws IOException { - throw new UnsupportedOperationException(WRITE_NOT_SUPPORTED); - } - - @Override - public Optional update( - org.hypertrace.core.documentstore.query.Query query, - java.util.Collection updates, - UpdateOptions updateOptions) - throws IOException { - throw new UnsupportedOperationException(WRITE_NOT_SUPPORTED); - } - - @Override - public CloseableIterator bulkUpdate( - org.hypertrace.core.documentstore.query.Query query, - java.util.Collection updates, - UpdateOptions updateOptions) - throws IOException { - throw new UnsupportedOperationException(WRITE_NOT_SUPPORTED); - } } From 57b623a5242fb828c9c95afa5b203351f322f6c0 Mon Sep 17 00:00:00 2001 From: Prashant Pandey Date: Tue, 13 Jan 2026 01:08:32 +0530 Subject: [PATCH 18/26] Implement create for flat collections --- .../FlatCollectionWriteTest.java | 139 ++-------- .../documentstore/commons/ColumnMetadata.java | 5 + .../postgres/FlatPostgresCollection.java | 254 +++++------------- .../postgres/PostgresMetadataFetcher.java | 8 +- .../model/PostgresColumnMetadata.java | 6 + .../postgres/PostgresMetadataFetcherTest.java | 64 +++++ 6 files changed, 174 insertions(+), 302 deletions(-) diff --git a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java index 68051082..269d926c 100644 --- a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java +++ b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java @@ -213,6 +213,15 @@ void testCreateNewDocument() throws Exception { objectNode.put("price", 999); objectNode.put("quantity", 50); objectNode.put("in_stock", true); + objectNode.set("tags", OBJECT_MAPPER.createArrayNode().add("electronics").add("sale")); + + // Add JSONB field + ObjectNode propsNode = OBJECT_MAPPER.createObjectNode(); + propsNode.put("color", "blue"); + propsNode.put("weight", 2.5); + propsNode.put("warranty", true); + objectNode.set("props", propsNode); + Document document = new JSONDocument(objectNode); Key key = new SingleValueKey("default", "new-doc-100"); @@ -233,6 +242,22 @@ void testCreateNewDocument() throws Exception { assertEquals(999, rs.getInt("price")); assertEquals(50, rs.getInt("quantity")); assertTrue(rs.getBoolean("in_stock")); + + // Verify tags array + java.sql.Array tagsArray = rs.getArray("tags"); + assertNotNull(tagsArray); + String[] tags = (String[]) tagsArray.getArray(); + assertEquals(2, tags.length); + assertEquals("electronics", tags[0]); + assertEquals("sale", tags[1]); + + // Verify JSONB props + String propsJson = rs.getString("props"); + assertNotNull(propsJson); + JsonNode propsResult = OBJECT_MAPPER.readTree(propsJson); + assertEquals("blue", propsResult.get("color").asText()); + assertEquals(2.5, propsResult.get("weight").asDouble(), 0.01); + assertTrue(propsResult.get("warranty").asBoolean()); } } @@ -315,120 +340,6 @@ void testCreateSkipsUnknownFields() throws Exception { assertEquals("Item", rs.getString("item")); } } - - @Test - @DisplayName("Should create a new document when key does not exist") - void testCreateOrReplaceNewDocument() throws Exception { - ObjectNode objectNode = OBJECT_MAPPER.createObjectNode(); - objectNode.put("id", "createorreplace-new-500"); - objectNode.put("item", "NewMirror"); - objectNode.put("price", 150); - Document document = new JSONDocument(objectNode); - Key key = new SingleValueKey("default", "createorreplace-new-500"); - - boolean result = flatCollection.createOrReplace(key, document); - - assertTrue(result); - - // Verify the data was inserted - PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; - try (Connection conn = pgDatastore.getPostgresClient(); - PreparedStatement ps = - conn.prepareStatement( - String.format( - "SELECT * FROM \"%s\" WHERE \"id\" = 'createorreplace-new-500'", - FLAT_COLLECTION_NAME)); - ResultSet rs = ps.executeQuery()) { - assertTrue(rs.next()); - assertEquals("NewMirror", rs.getString("item")); - assertEquals(150, rs.getInt("price")); - } - } - - @Test - @DisplayName("Should replace an existing document when key exists") - void testCreateOrReplaceExistingDocument() throws Exception { - // First create a document - ObjectNode objectNode1 = OBJECT_MAPPER.createObjectNode(); - objectNode1.put("id", "createorreplace-existing-600"); - objectNode1.put("item", "OriginalItem"); - objectNode1.put("price", 100); - Document document1 = new JSONDocument(objectNode1); - Key key = new SingleValueKey("default", "createorreplace-existing-600"); - flatCollection.create(key, document1); - - // Now replace it - ObjectNode objectNode2 = OBJECT_MAPPER.createObjectNode(); - objectNode2.put("id", "createorreplace-existing-600"); - objectNode2.put("item", "ReplacedItem"); - objectNode2.put("price", 200); - Document document2 = new JSONDocument(objectNode2); - - boolean result = flatCollection.createOrReplace(key, document2); - - assertTrue(result); - - // Verify the data was replaced - PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; - try (Connection conn = pgDatastore.getPostgresClient(); - PreparedStatement ps = - conn.prepareStatement( - String.format( - "SELECT * FROM \"%s\" WHERE \"id\" = 'createorreplace-existing-600'", - FLAT_COLLECTION_NAME)); - ResultSet rs = ps.executeQuery()) { - assertTrue(rs.next()); - assertEquals("ReplacedItem", rs.getString("item")); - assertEquals(200, rs.getInt("price")); - } - } - - @Test - @DisplayName("Should create and return a new document when key does not exist") - void testCreateOrReplaceAndReturnNewDocument() throws Exception { - ObjectNode objectNode = OBJECT_MAPPER.createObjectNode(); - objectNode.put("id", "createorreplace-return-700"); - objectNode.put("item", "ReturnedItem"); - objectNode.put("price", 250); - Document document = new JSONDocument(objectNode); - Key key = new SingleValueKey("default", "createorreplace-return-700"); - - Document result = flatCollection.createOrReplaceAndReturn(key, document); - - assertNotNull(result); - JsonNode resultNode = OBJECT_MAPPER.readTree(result.toJson()); - assertEquals("createorreplace-return-700", resultNode.get("id").asText()); - assertEquals("ReturnedItem", resultNode.get("item").asText()); - assertEquals(250, resultNode.get("price").asInt()); - } - - @Test - @DisplayName("Should replace and return an existing document when key exists") - void testCreateOrReplaceAndReturnExistingDocument() throws Exception { - // First create a document - ObjectNode objectNode1 = OBJECT_MAPPER.createObjectNode(); - objectNode1.put("id", "createorreplace-return-existing-800"); - objectNode1.put("item", "OriginalReturnItem"); - objectNode1.put("price", 300); - Document document1 = new JSONDocument(objectNode1); - Key key = new SingleValueKey("default", "createorreplace-return-existing-800"); - flatCollection.create(key, document1); - - // Now replace it - ObjectNode objectNode2 = OBJECT_MAPPER.createObjectNode(); - objectNode2.put("id", "createorreplace-return-existing-800"); - objectNode2.put("item", "ReplacedReturnItem"); - objectNode2.put("price", 400); - Document document2 = new JSONDocument(objectNode2); - - Document result = flatCollection.createOrReplaceAndReturn(key, document2); - - assertNotNull(result); - JsonNode resultNode = OBJECT_MAPPER.readTree(result.toJson()); - assertEquals("createorreplace-return-existing-800", resultNode.get("id").asText()); - assertEquals("ReplacedReturnItem", resultNode.get("item").asText()); - assertEquals(400, resultNode.get("price").asInt()); - } } @Nested diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/commons/ColumnMetadata.java b/document-store/src/main/java/org/hypertrace/core/documentstore/commons/ColumnMetadata.java index 475e55be..b52e16e9 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/commons/ColumnMetadata.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/commons/ColumnMetadata.java @@ -18,4 +18,9 @@ public interface ColumnMetadata { * @return whether this column can be set to NULL */ boolean isNullable(); + + /** + * @return whether this column is an array type + */ + boolean isArray(); } diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java index 6dedceef..422c9736 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java @@ -6,19 +6,18 @@ import java.sql.Connection; import java.sql.Date; import java.sql.PreparedStatement; -import java.sql.ResultSet; import java.sql.SQLException; import java.sql.Timestamp; import java.sql.Types; import java.time.Instant; import java.util.ArrayList; +import java.util.HashMap; import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.Map.Entry; import java.util.Optional; import java.util.Set; -import java.util.stream.Collectors; import org.hypertrace.core.documentstore.BulkArrayValueUpdateRequest; import org.hypertrace.core.documentstore.BulkDeleteResult; import org.hypertrace.core.documentstore.BulkUpdateRequest; @@ -27,7 +26,6 @@ import org.hypertrace.core.documentstore.CreateResult; import org.hypertrace.core.documentstore.Document; import org.hypertrace.core.documentstore.Filter; -import org.hypertrace.core.documentstore.JSONDocument; import org.hypertrace.core.documentstore.Key; import org.hypertrace.core.documentstore.UpdateResult; import org.hypertrace.core.documentstore.model.exception.DuplicateDocumentException; @@ -174,12 +172,12 @@ public CreateResult create(Key key, Document document) throws IOException { @Override public boolean createOrReplace(Key key, Document document) throws IOException { - return createOrReplaceWithRetry(key, document, false); + throw new UnsupportedOperationException(WRITE_NOT_SUPPORTED); } @Override public Document createOrReplaceAndReturn(Key key, Document document) throws IOException { - return createOrReplaceAndReturnWithRetry(key, document, false); + throw new UnsupportedOperationException(WRITE_NOT_SUPPORTED); } @Override @@ -240,60 +238,9 @@ private CreateResult createWithRetry(Key key, Document document, boolean isRetry } } - private boolean createOrReplaceWithRetry(Key key, Document document, boolean isRetry) - throws IOException { - String tableName = tableIdentifier.getTableName(); - - try { - TypedDocument parsed = parseDocument(document, tableName); - if (parsed.isEmpty()) { - LOGGER.warn("No valid columns found in the document for table: {}", tableName); - return false; - } - - String sql = buildUpsertSql(parsed.getColumns()); - LOGGER.debug("Upsert SQL: {}", sql); - - int result = executeUpdate(sql, parsed); - LOGGER.debug("CreateOrReplace result: {}", result); - return result > 0; - - } catch (PSQLException e) { - return handlePSQLExceptionForCreateOrReplace(e, key, document, tableName, isRetry); - } catch (SQLException e) { - LOGGER.error("SQLException in createOrReplace. key: {} content: {}", key, document, e); - throw new IOException(e); - } - } - - private Document createOrReplaceAndReturnWithRetry(Key key, Document document, boolean isRetry) - throws IOException { - String tableName = tableIdentifier.getTableName(); - - try { - TypedDocument parsed = parseDocument(document, tableName); - if (parsed.isEmpty()) { - LOGGER.warn("No valid columns found in the document for table: {}", tableName); - return null; - } - - String sql = buildUpsertSqlWithReturning(parsed.getColumns()); - LOGGER.debug("Upsert with RETURNING SQL: {}", sql); - - return executeQueryAndReturn(sql, parsed); - - } catch (PSQLException e) { - return handlePSQLExceptionForCreateOrReplaceAndReturn(e, key, document, tableName, isRetry); - } catch (SQLException e) { - LOGGER.error( - "SQLException in createOrReplaceAndReturn. key: {} content: {}", key, document, e); - throw new IOException(e); - } - } - private TypedDocument parseDocument(Document document, String tableName) throws IOException { JsonNode jsonNode = MAPPER.readTree(document.toJson()); - List entries = new ArrayList<>(); + TypedDocument typedDocument = new TypedDocument(); Iterator> fields = jsonNode.fields(); while (fields.hasNext()) { @@ -310,40 +257,31 @@ private TypedDocument parseDocument(Document document, String tableName) throws } PostgresDataType type = columnMetadata.get().getPostgresType(); - entries.add(new ColumnEntry("\"" + fieldName + "\"", extractValue(fieldValue, type), type)); + boolean isArray = columnMetadata.get().isArray(); + typedDocument.add( + "\"" + fieldName + "\"", extractValue(fieldValue, type, isArray), type, isArray); } - return new TypedDocument(entries); + return typedDocument; } private int executeUpdate(String sql, TypedDocument parsed) throws SQLException { try (Connection conn = client.getPooledConnection(); PreparedStatement ps = conn.prepareStatement(sql)) { int index = 1; - for (ColumnEntry entry : parsed.entries) { - setParameter(ps, index++, entry.value, entry.type); + for (String column : parsed.getColumns()) { + setParameter( + conn, + ps, + index++, + parsed.getValue(column), + parsed.getType(column), + parsed.isArray(column)); } return ps.executeUpdate(); } } - private Document executeQueryAndReturn(String sql, TypedDocument parsed) - throws SQLException, IOException { - try (Connection conn = client.getPooledConnection(); - PreparedStatement ps = conn.prepareStatement(sql)) { - int index = 1; - for (ColumnEntry entry : parsed.entries) { - setParameter(ps, index++, entry.value, entry.type); - } - try (ResultSet rs = ps.executeQuery()) { - if (rs.next()) { - return resultSetToDocument(rs); - } - return null; - } - } - } - private CreateResult handlePSQLExceptionForCreate( PSQLException e, Key key, Document document, String tableName, boolean isRetry) throws IOException { @@ -359,131 +297,45 @@ private CreateResult handlePSQLExceptionForCreate( throw new IOException(e); } - private boolean handlePSQLExceptionForCreateOrReplace( - PSQLException e, Key key, Document document, String tableName, boolean isRetry) - throws IOException { - if (!isRetry && shouldRefreshSchemaAndRetry(e.getSQLState())) { - LOGGER.info( - "Schema mismatch detected (SQLState: {}), refreshing schema and retrying. key: {}", - e.getSQLState(), - key); - schemaRegistry.invalidate(tableName); - return createOrReplaceWithRetry(key, document, true); - } - LOGGER.error("SQLException in createOrReplace. key: {} content: {}", key, document, e); - throw new IOException(e); - } - - private Document handlePSQLExceptionForCreateOrReplaceAndReturn( - PSQLException e, Key key, Document document, String tableName, boolean isRetry) - throws IOException { - if (!isRetry && shouldRefreshSchemaAndRetry(e.getSQLState())) { - LOGGER.info( - "Schema mismatch detected (SQLState: {}), refreshing schema and retrying. key: {}", - e.getSQLState(), - key); - schemaRegistry.invalidate(tableName); - return createOrReplaceAndReturnWithRetry(key, document, true); - } - LOGGER.error("SQLException in createOrReplaceAndReturn. key: {} content: {}", key, document, e); - throw new IOException(e); - } - private boolean shouldRefreshSchemaAndRetry(String sqlState) { return PSQLState.UNDEFINED_COLUMN.getState().equals(sqlState) || PSQLState.DATATYPE_MISMATCH.getState().equals(sqlState); } - private static class ColumnEntry { - final String column; - final Object value; - final PostgresDataType type; - - ColumnEntry(String column, Object value, PostgresDataType type) { - this.column = column; - this.value = value; - this.type = type; - } - } - + /** + * Typed document contains field information along with the field type. Uses LinkedHashMaps keyed + * by column name. LinkedHashMap preserves insertion order for consistent parameter binding. + */ private static class TypedDocument { - final List entries; - - TypedDocument(List entries) { - this.entries = entries; + private final Map values = new HashMap<>(); + private final Map types = new HashMap<>(); + private final Map arrays = new HashMap<>(); + + void add(String column, Object value, PostgresDataType type, boolean isArray) { + values.put(column, value); + types.put(column, type); + arrays.put(column, isArray); } boolean isEmpty() { - return entries.isEmpty(); + return values.isEmpty(); } List getColumns() { - return entries.stream().map(e -> e.column).collect(Collectors.toList()); + return new ArrayList<>(values.keySet()); } - } - private String buildUpsertSql(List columns) { - String columnList = String.join(", ", columns); - String placeholders = String.join(", ", columns.stream().map(c -> "?").toArray(String[]::new)); - String updateSet = - columns.stream() - .filter(col -> !"\"id\"".equals(col)) - .map(col -> col + " = EXCLUDED." + col) - .collect(Collectors.joining(", ")); - - if (updateSet.isEmpty()) { - return String.format( - "INSERT INTO %s (%s) VALUES (%s) ON CONFLICT (\"id\") DO NOTHING", - tableIdentifier, columnList, placeholders); + Object getValue(String column) { + return values.get(column); } - return String.format( - "INSERT INTO %s (%s) VALUES (%s) ON CONFLICT (\"id\") DO UPDATE SET %s", - tableIdentifier, columnList, placeholders, updateSet); - } - - private String buildUpsertSqlWithReturning(List columns) { - return buildUpsertSql(columns) + " RETURNING *"; - } - - private Document resultSetToDocument(ResultSet rs) throws SQLException, IOException { - int columnCount = rs.getMetaData().getColumnCount(); - com.fasterxml.jackson.databind.node.ObjectNode objectNode = MAPPER.createObjectNode(); - - for (int i = 1; i <= columnCount; i++) { - String columnName = rs.getMetaData().getColumnName(i); - Object value = rs.getObject(i); - - if (value == null) { - objectNode.putNull(columnName); - } else if (value instanceof Integer) { - objectNode.put(columnName, (Integer) value); - } else if (value instanceof Long) { - objectNode.put(columnName, (Long) value); - } else if (value instanceof Double) { - objectNode.put(columnName, (Double) value); - } else if (value instanceof Float) { - objectNode.put(columnName, (Float) value); - } else if (value instanceof Boolean) { - objectNode.put(columnName, (Boolean) value); - } else if (value instanceof Timestamp) { - objectNode.put(columnName, ((Timestamp) value).toInstant().toString()); - } else if (value instanceof java.sql.Date) { - objectNode.put(columnName, value.toString()); - } else if (value instanceof org.postgresql.util.PGobject) { - // Handle JSONB - String jsonValue = ((org.postgresql.util.PGobject) value).getValue(); - if (jsonValue != null) { - objectNode.set(columnName, MAPPER.readTree(jsonValue)); - } else { - objectNode.putNull(columnName); - } - } else { - objectNode.put(columnName, value.toString()); - } + PostgresDataType getType(String column) { + return types.get(column); } - return new JSONDocument(objectNode); + boolean isArray(String column) { + return arrays.getOrDefault(column, false); + } } private String buildInsertSql(List columns) { @@ -493,11 +345,26 @@ private String buildInsertSql(List columns) { "INSERT INTO %s (%s) VALUES (%s)", tableIdentifier, columnList, placeholders); } - private Object extractValue(JsonNode node, PostgresDataType type) { + private Object extractValue(JsonNode node, PostgresDataType type, boolean isArray) { if (node == null || node.isNull()) { return null; } + if (isArray) { + if (!node.isArray()) { + node = MAPPER.createArrayNode().add(node); + } + List values = new ArrayList<>(); + for (JsonNode element : node) { + values.add(extractScalarValue(element, type)); + } + return values.toArray(); + } + + return extractScalarValue(node, type); + } + + private Object extractScalarValue(JsonNode node, PostgresDataType type) { switch (type) { case INTEGER: return node.isNumber() ? node.intValue() : Integer.parseInt(node.asText()); @@ -528,13 +395,26 @@ private Object extractValue(JsonNode node, PostgresDataType type) { } } - private void setParameter(PreparedStatement ps, int index, Object value, PostgresDataType type) + private void setParameter( + Connection conn, + PreparedStatement ps, + int index, + Object value, + PostgresDataType type, + boolean isArray) throws SQLException { if (value == null) { ps.setObject(index, null); return; } + if (isArray) { + Object[] arrayValues = (value instanceof Object[]) ? (Object[]) value : new Object[] {value}; + java.sql.Array sqlArray = conn.createArrayOf(type.getSqlType(), arrayValues); + ps.setArray(index, sqlArray); + return; + } + switch (type) { case INTEGER: ps.setInt(index, (Integer) value); diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresMetadataFetcher.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresMetadataFetcher.java index 24b10e07..8e35937b 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresMetadataFetcher.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresMetadataFetcher.java @@ -37,10 +37,16 @@ public Map fetch(String tableName) { String columnName = rs.getString("column_name"); String udtName = rs.getString("udt_name"); boolean isNullable = "YES".equalsIgnoreCase(rs.getString("is_nullable")); + boolean isArray = udtName != null && udtName.startsWith("_"); + String baseType = isArray ? udtName.substring(1) : udtName; metadataMap.put( columnName, new PostgresColumnMetadata( - columnName, mapToCanonicalType(udtName), mapToPostgresType(udtName), isNullable)); + columnName, + mapToCanonicalType(baseType), + mapToPostgresType(baseType), + isNullable, + isArray)); } } return metadataMap; diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/model/PostgresColumnMetadata.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/model/PostgresColumnMetadata.java index 3a2d4540..8998d2e5 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/model/PostgresColumnMetadata.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/model/PostgresColumnMetadata.java @@ -15,6 +15,7 @@ public class PostgresColumnMetadata implements ColumnMetadata { private final DataType canonicalType; @Getter private final PostgresDataType postgresType; private final boolean nullable; + private final boolean array; @Override public String getName() { @@ -30,4 +31,9 @@ public DataType getCanonicalType() { public boolean isNullable() { return nullable; } + + @Override + public boolean isArray() { + return array; + } } diff --git a/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresMetadataFetcherTest.java b/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresMetadataFetcherTest.java index 6fa32d91..7c373eb8 100644 --- a/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresMetadataFetcherTest.java +++ b/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresMetadataFetcherTest.java @@ -300,6 +300,70 @@ void fetchThrowsRuntimeExceptionOnSqlException() throws SQLException { assertTrue(exception.getCause() instanceof SQLException); } + @Test + void fetchMapsTextArrayToStringArray() throws SQLException { + setupSingleColumnResult("col", "_text", "NO"); + + PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col"); + + assertEquals(DataType.STRING, meta.getCanonicalType()); + assertEquals(PostgresDataType.TEXT, meta.getPostgresType()); + assertTrue(meta.isArray()); + } + + @Test + void fetchMapsInt4ArrayToIntegerArray() throws SQLException { + setupSingleColumnResult("col", "_int4", "NO"); + + PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col"); + + assertEquals(DataType.INTEGER, meta.getCanonicalType()); + assertEquals(PostgresDataType.INTEGER, meta.getPostgresType()); + assertTrue(meta.isArray()); + } + + @Test + void fetchMapsInt8ArrayToLongArray() throws SQLException { + setupSingleColumnResult("col", "_int8", "NO"); + + PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col"); + + assertEquals(DataType.LONG, meta.getCanonicalType()); + assertEquals(PostgresDataType.BIGINT, meta.getPostgresType()); + assertTrue(meta.isArray()); + } + + @Test + void fetchMapsFloat8ArrayToDoubleArray() throws SQLException { + setupSingleColumnResult("col", "_float8", "NO"); + + PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col"); + + assertEquals(DataType.DOUBLE, meta.getCanonicalType()); + assertEquals(PostgresDataType.DOUBLE_PRECISION, meta.getPostgresType()); + assertTrue(meta.isArray()); + } + + @Test + void fetchMapsBoolArrayToBooleanArray() throws SQLException { + setupSingleColumnResult("col", "_bool", "NO"); + + PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col"); + + assertEquals(DataType.BOOLEAN, meta.getCanonicalType()); + assertEquals(PostgresDataType.BOOLEAN, meta.getPostgresType()); + assertTrue(meta.isArray()); + } + + @Test + void fetchReturnsIsArrayFalseForNonArrayTypes() throws SQLException { + setupSingleColumnResult("col", "text", "NO"); + + PostgresColumnMetadata meta = fetcher.fetch(TEST_TABLE).get("col"); + + assertFalse(meta.isArray()); + } + private void setupSingleColumnResult(String colName, String udtName, String isNullable) throws SQLException { when(resultSet.next()).thenReturn(true, false); From 9f8811e3b7b04321b273c80818f08cf7b2b53feb Mon Sep 17 00:00:00 2001 From: Prashant Pandey Date: Tue, 13 Jan 2026 01:41:46 +0530 Subject: [PATCH 19/26] Fix compilation issue --- .../core/documentstore/model/config/ConnectionConfig.java | 2 -- 1 file changed, 2 deletions(-) diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/model/config/ConnectionConfig.java b/document-store/src/main/java/org/hypertrace/core/documentstore/model/config/ConnectionConfig.java index deae97ef..11127b81 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/model/config/ConnectionConfig.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/model/config/ConnectionConfig.java @@ -127,8 +127,6 @@ public ConnectionConfig build() { applicationName, connectionPoolConfig, queryTimeout, - null, - null, customParameters); } From bfd665135ab7b6274ef897b66f072dab2609c2fb Mon Sep 17 00:00:00 2001 From: Prashant Pandey Date: Wed, 14 Jan 2026 11:52:18 +0530 Subject: [PATCH 20/26] Refactor --- .../documentstore/postgres/model/PostgresColumnMetadata.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/model/PostgresColumnMetadata.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/model/PostgresColumnMetadata.java index 8998d2e5..ecc6ca6a 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/model/PostgresColumnMetadata.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/model/PostgresColumnMetadata.java @@ -15,7 +15,7 @@ public class PostgresColumnMetadata implements ColumnMetadata { private final DataType canonicalType; @Getter private final PostgresDataType postgresType; private final boolean nullable; - private final boolean array; + private final boolean isArray; @Override public String getName() { @@ -34,6 +34,6 @@ public boolean isNullable() { @Override public boolean isArray() { - return array; + return isArray; } } From 70ec4b3f7d578e383cadf674af473a109b258f48 Mon Sep 17 00:00:00 2001 From: Prashant Pandey Date: Wed, 14 Jan 2026 12:11:36 +0530 Subject: [PATCH 21/26] Enhanced CreateResult.java and others --- .../FlatCollectionWriteTest.java | 57 +++++++++++++++++++ .../core/documentstore/CreateResult.java | 18 ++++-- .../TypesafeDatastoreConfigAdapter.java | 2 - .../postgres/FlatPostgresCollection.java | 17 ++++-- 4 files changed, 83 insertions(+), 11 deletions(-) diff --git a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java index 269d926c..c932171d 100644 --- a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java +++ b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java @@ -2,6 +2,7 @@ import static org.hypertrace.core.documentstore.utils.Utils.readFileFromResource; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -340,6 +341,62 @@ void testCreateSkipsUnknownFields() throws Exception { assertEquals("Item", rs.getString("item")); } } + + @Test + @DisplayName("Should return skipped fields in CreateResult when columns are missing") + void testCreateReturnsSkippedFieldsInResult() throws Exception { + ObjectNode objectNode = OBJECT_MAPPER.createObjectNode(); + objectNode.put("id", "skipped-fields-doc-500"); + objectNode.put("item", "Valid Item"); + objectNode.put("price", 100); + objectNode.put("nonexistent_field1", "value1"); + objectNode.put("nonexistent_field2", "value2"); + Document document = new JSONDocument(objectNode); + Key key = new SingleValueKey("default", "skipped-fields-doc-500"); + + CreateResult result = flatCollection.create(key, document); + + assertTrue(result.isSucceed()); + assertTrue(result.isPartial()); + assertNotNull(result.getSkippedFields()); + assertEquals(2, result.getSkippedFields().size()); + assertTrue( + result + .getSkippedFields() + .containsAll(List.of("nonexistent_field1", "nonexistent_field2"))); + + // Verify the valid fields were inserted + PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; + try (Connection conn = pgDatastore.getPostgresClient(); + PreparedStatement ps = + conn.prepareStatement( + String.format( + "SELECT * FROM \"%s\" WHERE \"id\" = 'skipped-fields-doc-500'", + FLAT_COLLECTION_NAME)); + ResultSet rs = ps.executeQuery()) { + assertTrue(rs.next()); + assertEquals("Valid Item", rs.getString("item")); + assertEquals(100, rs.getInt("price")); + } + } + + @Test + @DisplayName("Should return empty skipped fields when all columns exist") + void testCreateReturnsEmptySkippedFieldsWhenAllColumnsExist() throws Exception { + ObjectNode objectNode = OBJECT_MAPPER.createObjectNode(); + objectNode.put("id", "all-valid-doc-600"); + objectNode.put("item", "Valid Item"); + objectNode.put("price", 200); + objectNode.put("quantity", 10); + Document document = new JSONDocument(objectNode); + Key key = new SingleValueKey("default", "all-valid-doc-600"); + + CreateResult result = flatCollection.create(key, document); + + assertTrue(result.isSucceed()); + assertFalse(result.isPartial()); + assertTrue(result.getSkippedFields().isEmpty()); + } } @Nested diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/CreateResult.java b/document-store/src/main/java/org/hypertrace/core/documentstore/CreateResult.java index 38934c7f..7b0f7736 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/CreateResult.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/CreateResult.java @@ -1,16 +1,24 @@ package org.hypertrace.core.documentstore; +import java.util.List; +import lombok.AllArgsConstructor; +import lombok.Getter; + /* * Represent the result object for CREATE operation of document store APIs. * */ +@AllArgsConstructor +@Getter public class CreateResult { - private boolean succeed; + private boolean isSucceed; + private boolean onRetry; + private List skippedFields; - public CreateResult(boolean succeed) { - this.succeed = succeed; + public CreateResult(boolean isSucceed) { + this.isSucceed = isSucceed; } - public boolean isSucceed() { - return succeed; + public boolean isPartial() { + return !skippedFields.isEmpty(); } } diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/TypesafeDatastoreConfigAdapter.java b/document-store/src/main/java/org/hypertrace/core/documentstore/TypesafeDatastoreConfigAdapter.java index 4fc19e55..abc00e88 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/TypesafeDatastoreConfigAdapter.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/TypesafeDatastoreConfigAdapter.java @@ -94,8 +94,6 @@ public DatastoreConfig convert(final Config config) { connectionConfig.applicationName(), connectionConfig.connectionPoolConfig(), connectionConfig.queryTimeout(), - connectionConfig.schemaCacheExpiry(), - connectionConfig.schemaRefreshCooldown(), connectionConfig.customParameters()) { @Override public String toConnectionString() { diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java index 422c9736..3fa67769 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java @@ -208,16 +208,19 @@ public CloseableIterator bulkUpdate( throw new UnsupportedOperationException(WRITE_NOT_SUPPORTED); } + /*isRetry: Whether this is a retry attempt*/ private CreateResult createWithRetry(Key key, Document document, boolean isRetry) throws IOException { String tableName = tableIdentifier.getTableName(); + List skippedFields = new ArrayList<>(); + try { - TypedDocument parsed = parseDocument(document, tableName); + TypedDocument parsed = parseDocument(document, tableName, skippedFields); // if there are no valid columns in the document if (parsed.isEmpty()) { LOGGER.warn("No valid columns found in the document for table: {}", tableName); - return new CreateResult(false); + return new CreateResult(false, isRetry, skippedFields); } String sql = buildInsertSql(parsed.getColumns()); @@ -225,7 +228,7 @@ private CreateResult createWithRetry(Key key, Document document, boolean isRetry int result = executeUpdate(sql, parsed); LOGGER.debug("Create result: {}", result); - return new CreateResult(result > 0); + return new CreateResult(result > 0, isRetry, skippedFields); } catch (PSQLException e) { if (PSQLState.UNIQUE_VIOLATION.getState().equals(e.getSQLState())) { @@ -238,7 +241,8 @@ private CreateResult createWithRetry(Key key, Document document, boolean isRetry } } - private TypedDocument parseDocument(Document document, String tableName) throws IOException { + private TypedDocument parseDocument( + Document document, String tableName, List skippedColumns) throws IOException { JsonNode jsonNode = MAPPER.readTree(document.toJson()); TypedDocument typedDocument = new TypedDocument(); @@ -253,6 +257,7 @@ private TypedDocument parseDocument(Document document, String tableName) throws if (columnMetadata.isEmpty()) { LOGGER.warn("Could not find column metadata for column: {}, skipping it", fieldName); + skippedColumns.add(fieldName); continue; } @@ -297,6 +302,10 @@ private CreateResult handlePSQLExceptionForCreate( throw new IOException(e); } + /** + * Returns true if the SQL state indicates a schema mismatch, i.e. the column does not exist or + * the data type is mismatched. + */ private boolean shouldRefreshSchemaAndRetry(String sqlState) { return PSQLState.UNDEFINED_COLUMN.getState().equals(sqlState) || PSQLState.DATATYPE_MISMATCH.getState().equals(sqlState); From 6d1c277b5f09472863120992beb043f252cb8924 Mon Sep 17 00:00:00 2001 From: Prashant Pandey Date: Wed, 14 Jan 2026 12:34:21 +0530 Subject: [PATCH 22/26] Added more test cases --- .../FlatCollectionWriteTest.java | 142 ++++++++++++++++++ .../postgres/FlatPostgresCollection.java | 15 +- 2 files changed, 155 insertions(+), 2 deletions(-) diff --git a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java index c932171d..ea2336af 100644 --- a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java +++ b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java @@ -397,6 +397,148 @@ void testCreateReturnsEmptySkippedFieldsWhenAllColumnsExist() throws Exception { assertFalse(result.isPartial()); assertTrue(result.getSkippedFields().isEmpty()); } + + @Test + @DisplayName("Should return failure when all fields are unknown (parsed.isEmpty)") + void testCreateFailsWhenAllFieldsAreUnknown() throws Exception { + // Document with only unknown fields - no valid columns will be found + ObjectNode objectNode = OBJECT_MAPPER.createObjectNode(); + objectNode.put("completely_unknown_field1", "value1"); + objectNode.put("completely_unknown_field2", "value2"); + objectNode.put("another_nonexistent_column", 123); + Document document = new JSONDocument(objectNode); + Key key = new SingleValueKey("default", "all-unknown-doc-700"); + + CreateResult result = flatCollection.create(key, document); + + // Should fail because no valid columns found (parsed.isEmpty() == true) + assertFalse(result.isSucceed()); + assertEquals(3, result.getSkippedFields().size()); + assertTrue( + result + .getSkippedFields() + .containsAll( + List.of( + "completely_unknown_field1", + "completely_unknown_field2", + "another_nonexistent_column"))); + + // Verify no row was inserted + PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; + try (Connection conn = pgDatastore.getPostgresClient(); + PreparedStatement ps = + conn.prepareStatement( + String.format( + "SELECT COUNT(*) FROM \"%s\" WHERE \"id\" = 'all-unknown-doc-700'", + FLAT_COLLECTION_NAME)); + ResultSet rs = ps.executeQuery()) { + assertTrue(rs.next()); + assertEquals(0, rs.getInt(1)); + } + } + + @Test + @DisplayName("Should refresh schema and retry on UNDEFINED_COLUMN error") + void testCreateRefreshesSchemaOnUndefinedColumnError() throws Exception { + PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; + + // Step 1: Add a temporary column and do a create to cache the schema + String addColumnSQL = + String.format( + "ALTER TABLE \"%s\" ADD COLUMN \"temp_col\" TEXT", FLAT_COLLECTION_NAME); + try (Connection conn = pgDatastore.getPostgresClient(); + PreparedStatement ps = conn.prepareStatement(addColumnSQL)) { + ps.execute(); + LOGGER.info("Added temporary column 'temp_col' to table"); + } + + // Step 2: Create a document with the temp column to cache the schema + ObjectNode objectNode1 = OBJECT_MAPPER.createObjectNode(); + objectNode1.put("id", "cache-schema-doc"); + objectNode1.put("item", "Item to cache schema"); + objectNode1.put("temp_col", "temp value"); + flatCollection.create( + new SingleValueKey("default", "cache-schema-doc"), new JSONDocument(objectNode1)); + LOGGER.info("Schema cached with temp_col"); + + // Step 3: DROP the column - now the cached schema is stale + String dropColumnSQL = + String.format("ALTER TABLE \"%s\" DROP COLUMN \"temp_col\"", FLAT_COLLECTION_NAME); + try (Connection conn = pgDatastore.getPostgresClient(); + PreparedStatement ps = conn.prepareStatement(dropColumnSQL)) { + ps.execute(); + LOGGER.info("Dropped temp_col - schema cache is now stale"); + } + + // Step 4: Try to create with the dropped column + // Schema registry still thinks temp_col exists, so it will include it in INSERT + // INSERT will fail with UNDEFINED_COLUMN, triggering handlePSQLExceptionForCreate + // which will refresh schema and retry + ObjectNode objectNode2 = OBJECT_MAPPER.createObjectNode(); + objectNode2.put("id", "retry-doc-800"); + objectNode2.put("item", "Item after schema refresh"); + objectNode2.put("temp_col", "this column no longer exists"); + Document document = new JSONDocument(objectNode2); + Key key = new SingleValueKey("default", "retry-doc-800"); + + CreateResult result = flatCollection.create(key, document); + + // Should succeed after retry - temp_col will be skipped on retry + assertTrue(result.isSucceed()); + // On retry, the column won't be found and will be skipped + assertTrue(result.getSkippedFields().contains("temp_col")); + // Should be marked as retry + assertTrue(result.isOnRetry()); + + // Verify the valid fields were inserted + try (Connection conn = pgDatastore.getPostgresClient(); + PreparedStatement ps = + conn.prepareStatement( + String.format( + "SELECT * FROM \"%s\" WHERE \"id\" = 'retry-doc-800'", + FLAT_COLLECTION_NAME)); + ResultSet rs = ps.executeQuery()) { + assertTrue(rs.next()); + assertEquals("Item after schema refresh", rs.getString("item")); + } + } + + @Test + @DisplayName("Should skip column with unparseable value and add to skippedFields") + void testCreateSkipsUnparseableValues() throws Exception { + // Try to insert a string value into an integer column with wrong type + // The unparseable column should be skipped, not throw an exception + ObjectNode objectNode = OBJECT_MAPPER.createObjectNode(); + objectNode.put("id", "datatype-mismatch-doc-900"); + objectNode.put("item", "Valid Item"); + objectNode.put("price", "not_a_number_at_all"); // price is INTEGER, this will fail parsing + Document document = new JSONDocument(objectNode); + Key key = new SingleValueKey("default", "datatype-mismatch-doc-900"); + + CreateResult result = flatCollection.create(key, document); + + // Should succeed with the valid columns, skipping the unparseable one + assertTrue(result.isSucceed()); + assertTrue(result.isPartial()); + assertEquals(1, result.getSkippedFields().size()); + assertTrue(result.getSkippedFields().contains("price")); + + // Verify the valid fields were inserted + PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; + try (Connection conn = pgDatastore.getPostgresClient(); + PreparedStatement ps = + conn.prepareStatement( + String.format( + "SELECT * FROM \"%s\" WHERE \"id\" = 'datatype-mismatch-doc-900'", + FLAT_COLLECTION_NAME)); + ResultSet rs = ps.executeQuery()) { + assertTrue(rs.next()); + assertEquals("Valid Item", rs.getString("item")); + // price should be null since it was skipped + assertEquals(0, rs.getInt("price")); + assertTrue(rs.wasNull()); + } + } } @Nested diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java index 3fa67769..04a5f901 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java @@ -263,8 +263,19 @@ private TypedDocument parseDocument( PostgresDataType type = columnMetadata.get().getPostgresType(); boolean isArray = columnMetadata.get().isArray(); - typedDocument.add( - "\"" + fieldName + "\"", extractValue(fieldValue, type, isArray), type, isArray); + + try { + Object value = extractValue(fieldValue, type, isArray); + typedDocument.add("\"" + fieldName + "\"", value, type, isArray); + } catch (Exception e) { + //If we fail to parse the value, we skip this field to write on a best-effort basis + LOGGER.warn( + "Could not parse value for column: {} with type: {}, skipping it. Error: {}", + fieldName, + type, + e.getMessage()); + skippedColumns.add(fieldName); + } } return typedDocument; From 910ef8c97ab9b34fdf3723f40493de9f1cd18317 Mon Sep 17 00:00:00 2001 From: Prashant Pandey Date: Wed, 14 Jan 2026 12:37:19 +0530 Subject: [PATCH 23/26] Spotless --- .../hypertrace/core/documentstore/FlatCollectionWriteTest.java | 3 +-- .../core/documentstore/postgres/FlatPostgresCollection.java | 2 +- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java index ea2336af..6979ccc1 100644 --- a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java +++ b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java @@ -444,8 +444,7 @@ void testCreateRefreshesSchemaOnUndefinedColumnError() throws Exception { // Step 1: Add a temporary column and do a create to cache the schema String addColumnSQL = - String.format( - "ALTER TABLE \"%s\" ADD COLUMN \"temp_col\" TEXT", FLAT_COLLECTION_NAME); + String.format("ALTER TABLE \"%s\" ADD COLUMN \"temp_col\" TEXT", FLAT_COLLECTION_NAME); try (Connection conn = pgDatastore.getPostgresClient(); PreparedStatement ps = conn.prepareStatement(addColumnSQL)) { ps.execute(); diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java index 04a5f901..071a259d 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java @@ -268,7 +268,7 @@ private TypedDocument parseDocument( Object value = extractValue(fieldValue, type, isArray); typedDocument.add("\"" + fieldName + "\"", value, type, isArray); } catch (Exception e) { - //If we fail to parse the value, we skip this field to write on a best-effort basis + // If we fail to parse the value, we skip this field to write on a best-effort basis LOGGER.warn( "Could not parse value for column: {} with type: {}, skipping it. Error: {}", fieldName, From 3e2c178507bb8124f5c6c9a56050f04d7bc1aa43 Mon Sep 17 00:00:00 2001 From: Prashant Pandey Date: Wed, 14 Jan 2026 13:47:22 +0530 Subject: [PATCH 24/26] WIP --- .../core/documentstore/FlatCollectionWriteTest.java | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java index 6979ccc1..bc5b1299 100644 --- a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java +++ b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java @@ -482,12 +482,10 @@ void testCreateRefreshesSchemaOnUndefinedColumnError() throws Exception { CreateResult result = flatCollection.create(key, document); - // Should succeed after retry - temp_col will be skipped on retry + // Should succeed - temp_col will be skipped (either via retry or schema refresh) assertTrue(result.isSucceed()); - // On retry, the column won't be found and will be skipped + // The dropped column should be skipped assertTrue(result.getSkippedFields().contains("temp_col")); - // Should be marked as retry - assertTrue(result.isOnRetry()); // Verify the valid fields were inserted try (Connection conn = pgDatastore.getPostgresClient(); From 9061e24fa7fd1aea9c41d4738a44e3a815dc99ef Mon Sep 17 00:00:00 2001 From: Prashant Pandey Date: Wed, 14 Jan 2026 14:00:44 +0530 Subject: [PATCH 25/26] Add more test coverage --- .../FlatCollectionWriteTest.java | 155 +++++++++++++++++- 1 file changed, 154 insertions(+), 1 deletion(-) diff --git a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java index bc5b1299..d109d637 100644 --- a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java +++ b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java @@ -99,7 +99,11 @@ private static void createFlatCollectionSchema() { + "\"sales\" JSONB," + "\"numbers\" INTEGER[]," + "\"scores\" DOUBLE PRECISION[]," - + "\"flags\" BOOLEAN[]" + + "\"flags\" BOOLEAN[]," + + "\"big_number\" BIGINT," + + "\"rating\" REAL," + + "\"created_date\" DATE," + + "\"weight\" DOUBLE PRECISION" + ");", FLAT_COLLECTION_NAME); @@ -536,6 +540,155 @@ void testCreateSkipsUnparseableValues() throws Exception { assertTrue(rs.wasNull()); } } + + @Test + @DisplayName("Should handle all scalar data types including string parsing and nulls") + void testCreateWithAllDataTypes() throws Exception { + // Test 1: All types with native values (number nodes, boolean nodes, etc.) + ObjectNode nativeTypesNode = OBJECT_MAPPER.createObjectNode(); + nativeTypesNode.put("id", "native-types-doc"); + nativeTypesNode.put("item", "Native Types"); // TEXT + nativeTypesNode.put("price", 100); // INTEGER (number node) + nativeTypesNode.put("big_number", 9223372036854775807L); // BIGINT (number node) + nativeTypesNode.put("rating", 4.5f); // REAL (number node) + nativeTypesNode.put("weight", 123.456789); // DOUBLE PRECISION (number node) + nativeTypesNode.put("in_stock", true); // BOOLEAN (boolean node) + nativeTypesNode.put("date", "2024-01-15T10:30:00Z"); // TIMESTAMPTZ (textual) + nativeTypesNode.put("created_date", "2024-01-15"); // DATE (textual) + nativeTypesNode.putObject("props").put("key", "value"); // JSONB + + CreateResult result1 = + flatCollection.create( + new SingleValueKey("default", "native-types-doc"), new JSONDocument(nativeTypesNode)); + assertTrue(result1.isSucceed()); + + // Test 2: String representations of numbers (covers parseInt, parseLong, etc.) + ObjectNode stringTypesNode = OBJECT_MAPPER.createObjectNode(); + stringTypesNode.put("id", "string-types-doc"); + stringTypesNode.put("item", "String Types"); + stringTypesNode.put("price", "200"); // INTEGER from string + stringTypesNode.put("big_number", "1234567890123"); // BIGINT from string + stringTypesNode.put("rating", "3.75"); // REAL from string + stringTypesNode.put("weight", "987.654"); // DOUBLE PRECISION from string + stringTypesNode.put("in_stock", "true"); // BOOLEAN from string + + CreateResult result2 = + flatCollection.create( + new SingleValueKey("default", "string-types-doc"), new JSONDocument(stringTypesNode)); + assertTrue(result2.isSucceed()); + + // Test 3: TIMESTAMPTZ from epoch milliseconds + long epochMillis = 1705315800000L; + ObjectNode epochNode = OBJECT_MAPPER.createObjectNode(); + epochNode.put("id", "epoch-doc"); + epochNode.put("item", "Epoch Timestamp"); + epochNode.put("date", epochMillis); // TIMESTAMPTZ from number + + CreateResult result3 = + flatCollection.create( + new SingleValueKey("default", "epoch-doc"), new JSONDocument(epochNode)); + assertTrue(result3.isSucceed()); + + // Test 4: Null values (covers setParameter null handling) + ObjectNode nullNode = OBJECT_MAPPER.createObjectNode(); + nullNode.put("id", "null-doc"); + nullNode.put("item", "Null Values"); + nullNode.putNull("price"); + nullNode.putNull("date"); + nullNode.putNull("in_stock"); + + CreateResult result4 = + flatCollection.create( + new SingleValueKey("default", "null-doc"), new JSONDocument(nullNode)); + assertTrue(result4.isSucceed()); + + // Verify all inserts + PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; + try (Connection conn = pgDatastore.getPostgresClient(); + PreparedStatement ps = + conn.prepareStatement( + String.format( + "SELECT * FROM \"%s\" WHERE \"id\" IN ('native-types-doc', 'string-types-doc', 'epoch-doc', 'null-doc') ORDER BY \"id\"", + FLAT_COLLECTION_NAME)); + ResultSet rs = ps.executeQuery()) { + + // epoch-doc + assertTrue(rs.next()); + assertEquals(epochMillis, rs.getTimestamp("date").getTime()); + + // native-types-doc + assertTrue(rs.next()); + assertEquals(100, rs.getInt("price")); + assertEquals(9223372036854775807L, rs.getLong("big_number")); + assertEquals(4.5f, rs.getFloat("rating"), 0.01f); + assertEquals(123.456789, rs.getDouble("weight"), 0.0001); + assertTrue(rs.getBoolean("in_stock")); + + // null-doc + assertTrue(rs.next()); + rs.getInt("price"); + assertTrue(rs.wasNull()); + + // string-types-doc + assertTrue(rs.next()); + assertEquals(200, rs.getInt("price")); + assertEquals(1234567890123L, rs.getLong("big_number")); + assertEquals(3.75f, rs.getFloat("rating"), 0.01f); + } + } + + @Test + @DisplayName("Should handle array types and single-to-array conversion") + void testCreateWithArrayTypes() throws Exception { + PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; + + // Test 1: Proper arrays + ObjectNode arrayNode = OBJECT_MAPPER.createObjectNode(); + arrayNode.put("id", "array-doc"); + arrayNode.put("item", "Array Types"); + arrayNode.putArray("tags").add("tag1").add("tag2"); // TEXT[] + arrayNode.putArray("numbers").add(10).add(20); // INTEGER[] + arrayNode.putArray("scores").add(1.5).add(2.5); // DOUBLE PRECISION[] + arrayNode.putArray("flags").add(true).add(false); // BOOLEAN[] + + CreateResult result1 = + flatCollection.create( + new SingleValueKey("default", "array-doc"), new JSONDocument(arrayNode)); + assertTrue(result1.isSucceed()); + + // Test 2: Single values auto-converted to arrays + ObjectNode singleNode = OBJECT_MAPPER.createObjectNode(); + singleNode.put("id", "single-to-array-doc"); + singleNode.put("item", "Single to Array"); + singleNode.put("tags", "single-tag"); // TEXT[] from single value + singleNode.put("numbers", 42); // INTEGER[] from single value + + CreateResult result2 = + flatCollection.create( + new SingleValueKey("default", "single-to-array-doc"), new JSONDocument(singleNode)); + assertTrue(result2.isSucceed()); + + // Verify + try (Connection conn = pgDatastore.getPostgresClient(); + PreparedStatement ps = + conn.prepareStatement( + String.format( + "SELECT * FROM \"%s\" WHERE \"id\" IN ('array-doc', 'single-to-array-doc') ORDER BY \"id\"", + FLAT_COLLECTION_NAME)); + ResultSet rs = ps.executeQuery()) { + + // array-doc + assertTrue(rs.next()); + assertEquals(2, ((String[]) rs.getArray("tags").getArray()).length); + assertEquals(2, ((Integer[]) rs.getArray("numbers").getArray()).length); + + // single-to-array-doc + assertTrue(rs.next()); + String[] tags = (String[]) rs.getArray("tags").getArray(); + assertEquals(1, tags.length); + assertEquals("single-tag", tags[0]); + } + } } @Nested From c3024b008d7ca75c5b8ff94d511efd42152acee8 Mon Sep 17 00:00:00 2001 From: Prashant Pandey Date: Fri, 16 Jan 2026 12:44:34 +0530 Subject: [PATCH 26/26] Added `bestEfforts` configuration to PG custom parameters. --- .../FlatCollectionWriteTest.java | 113 ++++++++++++++++++ .../exception/SchemaMismatchException.java | 22 ++++ .../postgres/FlatPostgresCollection.java | 27 ++++- 3 files changed, 161 insertions(+), 1 deletion(-) create mode 100644 document-store/src/main/java/org/hypertrace/core/documentstore/model/exception/SchemaMismatchException.java diff --git a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java index d109d637..cbe680ce 100644 --- a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java +++ b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java @@ -20,6 +20,7 @@ import java.util.Map; import java.util.Set; import org.hypertrace.core.documentstore.model.exception.DuplicateDocumentException; +import org.hypertrace.core.documentstore.model.exception.SchemaMismatchException; import org.hypertrace.core.documentstore.postgres.PostgresDatastore; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.AfterEach; @@ -848,4 +849,116 @@ void testBulkOperationOnArrayValue() throws IOException { () -> flatCollection.bulkOperationOnArrayValue(request)); } } + + @Nested + @DisplayName("Strict Mode (bestEffortWrites=false)") + class StrictModeTests { + + private Datastore strictDatastore; + private Collection strictCollection; + + @BeforeEach + void setupStrictModeDatastore() { + // Create a datastore with bestEffortWrites=false (strict mode) + String postgresConnectionUrl = + String.format("jdbc:postgresql://localhost:%s/", postgres.getMappedPort(5432)); + + Map strictConfig = new HashMap<>(); + strictConfig.put("url", postgresConnectionUrl); + strictConfig.put("user", "postgres"); + strictConfig.put("password", "postgres"); + // Configure strict mode via customParams + Map customParams = new HashMap<>(); + customParams.put("bestEffortWrites", "false"); + strictConfig.put("customParams", customParams); + + strictDatastore = + DatastoreProvider.getDatastore("Postgres", ConfigFactory.parseMap(strictConfig)); + strictCollection = + strictDatastore.getCollectionForType(FLAT_COLLECTION_NAME, DocumentType.FLAT); + } + + @Test + @DisplayName("Should throw SchemaMismatchException when column not in schema (strict mode)") + void testStrictModeThrowsOnUnknownColumn() { + ObjectNode objectNode = OBJECT_MAPPER.createObjectNode(); + objectNode.put("id", "strict-unknown-col-doc"); + objectNode.put("item", "Valid Item"); + objectNode.put("unknown_column", "this should fail"); + Document document = new JSONDocument(objectNode); + Key key = new SingleValueKey("default", "strict-unknown-col-doc"); + + SchemaMismatchException exception = + assertThrows(SchemaMismatchException.class, () -> strictCollection.create(key, document)); + + assertTrue(exception.getMessage().contains("unknown_column")); + assertTrue(exception.getMessage().contains("not found in schema")); + } + + @Test + @DisplayName( + "Should throw SchemaMismatchException when value type doesn't match schema (strict mode)") + void testStrictModeThrowsOnTypeMismatch() { + ObjectNode objectNode = OBJECT_MAPPER.createObjectNode(); + objectNode.put("id", "strict-type-mismatch-doc"); + objectNode.put("item", "Valid Item"); + objectNode.put("price", "not_a_number_at_all"); // price is INTEGER + Document document = new JSONDocument(objectNode); + Key key = new SingleValueKey("default", "strict-type-mismatch-doc"); + + SchemaMismatchException exception = + assertThrows(SchemaMismatchException.class, () -> strictCollection.create(key, document)); + + assertTrue(exception.getMessage().contains("price")); + assertTrue(exception.getMessage().contains("Failed to parse value")); + } + + @Test + @DisplayName("Should succeed in strict mode when all fields match schema") + void testStrictModeSucceedsWithValidDocument() throws Exception { + ObjectNode objectNode = OBJECT_MAPPER.createObjectNode(); + objectNode.put("id", "strict-valid-doc"); + objectNode.put("item", "Valid Item"); + objectNode.put("price", 100); + objectNode.put("quantity", 5); + objectNode.put("in_stock", true); + Document document = new JSONDocument(objectNode); + Key key = new SingleValueKey("default", "strict-valid-doc"); + + CreateResult result = strictCollection.create(key, document); + + assertTrue(result.isSucceed()); + assertFalse(result.isPartial()); + assertTrue(result.getSkippedFields().isEmpty()); + + // Verify data was inserted + PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; + try (Connection conn = pgDatastore.getPostgresClient(); + PreparedStatement ps = + conn.prepareStatement( + String.format( + "SELECT * FROM \"%s\" WHERE \"id\" = 'strict-valid-doc'", + FLAT_COLLECTION_NAME)); + ResultSet rs = ps.executeQuery()) { + assertTrue(rs.next()); + assertEquals("Valid Item", rs.getString("item")); + assertEquals(100, rs.getInt("price")); + } + } + + @Test + @DisplayName("Should throw SchemaMismatchException on first unknown field (strict mode)") + void testStrictModeFailsFastOnFirstUnknownField() { + ObjectNode objectNode = OBJECT_MAPPER.createObjectNode(); + objectNode.put("id", "strict-multi-unknown-doc"); + objectNode.put("unknown_field_1", "value1"); + objectNode.put("unknown_field_2", "value2"); + objectNode.put("item", "Valid Item"); + Document document = new JSONDocument(objectNode); + Key key = new SingleValueKey("default", "strict-multi-unknown-doc"); + + // Should throw on the first unknown field encountered + assertThrows(SchemaMismatchException.class, () -> strictCollection.create(key, document)); + } + } } diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/model/exception/SchemaMismatchException.java b/document-store/src/main/java/org/hypertrace/core/documentstore/model/exception/SchemaMismatchException.java new file mode 100644 index 00000000..b7f3afca --- /dev/null +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/model/exception/SchemaMismatchException.java @@ -0,0 +1,22 @@ +package org.hypertrace.core.documentstore.model.exception; + +import java.io.IOException; + +/** + * Exception thrown when a document field doesn't match the expected schema. This can occur when: + * + *
    + *
  • A field in the document doesn't exist in the schema + *
  • A field's value type doesn't match the expected schema type + *
+ */ +public class SchemaMismatchException extends IOException { + + public SchemaMismatchException(String message) { + super(message); + } + + public SchemaMismatchException(String message, Throwable cause) { + super(message, cause); + } +} diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java index 071a259d..2ae0f914 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java @@ -29,6 +29,7 @@ import org.hypertrace.core.documentstore.Key; import org.hypertrace.core.documentstore.UpdateResult; import org.hypertrace.core.documentstore.model.exception.DuplicateDocumentException; +import org.hypertrace.core.documentstore.model.exception.SchemaMismatchException; import org.hypertrace.core.documentstore.model.options.QueryOptions; import org.hypertrace.core.documentstore.model.options.UpdateOptions; import org.hypertrace.core.documentstore.model.subdoc.SubDocumentUpdate; @@ -55,15 +56,26 @@ public class FlatPostgresCollection extends PostgresCollection { private static final ObjectMapper MAPPER = new ObjectMapper(); private static final String WRITE_NOT_SUPPORTED = "Write operations are not supported for flat collections yet!"; + private static final String BEST_EFFORT_WRITES_CONFIG = "bestEffortWrites"; private final PostgresLazyilyLoadedSchemaRegistry schemaRegistry; + /** + * When true (default), fields that don't match the schema are skipped. When false (strict mode), + * all fields must be present in the schema with correct types (any fields present in the doc that + * are not present in the schema would make this check fail) + */ + private final boolean bestEffortWrites; + FlatPostgresCollection( final PostgresClient client, final String collectionName, final PostgresLazyilyLoadedSchemaRegistry schemaRegistry) { super(client, collectionName); this.schemaRegistry = schemaRegistry; + this.bestEffortWrites = + Boolean.parseBoolean( + client.getCustomParameters().getOrDefault(BEST_EFFORT_WRITES_CONFIG, "true")); } @Override @@ -256,6 +268,10 @@ private TypedDocument parseDocument( schemaRegistry.getColumnOrRefresh(tableName, fieldName); if (columnMetadata.isEmpty()) { + if (!bestEffortWrites) { + throw new SchemaMismatchException( + "Column '" + fieldName + "' not found in schema for table: " + tableName); + } LOGGER.warn("Could not find column metadata for column: {}, skipping it", fieldName); skippedColumns.add(fieldName); continue; @@ -268,7 +284,16 @@ private TypedDocument parseDocument( Object value = extractValue(fieldValue, type, isArray); typedDocument.add("\"" + fieldName + "\"", value, type, isArray); } catch (Exception e) { - // If we fail to parse the value, we skip this field to write on a best-effort basis + if (!bestEffortWrites) { + throw new SchemaMismatchException( + "Failed to parse value for column '" + + fieldName + + "' with type " + + type + + ": " + + e.getMessage(), + e); + } LOGGER.warn( "Could not parse value for column: {} with type: {}, skipping it. Error: {}", fieldName,