WIP

Author: suddendust

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<String, PostgresColumnMetadata> 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<String> columns = new ArrayList<>();
-      List<Object> values = new ArrayList<>();
-
-      // Extract fields from document that exist in schema
-      Iterator<Map.Entry<String, JsonNode>> fields = docJson.fields();
-      while (fields.hasNext()) {
-        Map.Entry<String, JsonNode> 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<String> 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

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),

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<String, PostgresColumnMetadata> fetch(String tableName) {
     Map<String, PostgresColumnMetadata> metadataMap = new HashMap<>();
 
-    try (Connection conn = datastore.getPostgresClient();
+    try (Connection conn = client.getPooledConnection();
         PreparedStatement ps = conn.prepareStatement(DISCOVERY_SQL)) {
 
       ps.setString(1, tableName);