Refactor Postgres Field Transformers and Collections

document-store/src/integrationTest/resources/query/pg_flat_collection_insert.json

@@ -0,0 +1,12 @@
+{
+  "statements": [
+    "INSERT INTO \"myTestFlat\" (\n\"_id\", \"item\", \"price\", \"quantity\", \"date\", \"tags\", \"props\", \"sales\"\n) VALUES (\n1, 'Soap', 10, 2, '2014-03-01T08:00:00Z',\n'{\"hygiene\", \"personal-care\", \"premium\"}',\n'{\"colors\": [\"Blue\", \"Green\"], \"brand\": \"Dettol\", \"size\": \"M\", \"seller\": {\"name\": \"Metro Chemicals Pvt. Ltd.\", \"address\": {\"city\": \"Mumbai\", \"pincode\": 400004}}}',\nNULL\n)",
+    "INSERT INTO \"myTestFlat\" (\n\"_id\", \"item\", \"price\", \"quantity\", \"date\", \"tags\", \"props\", \"sales\"\n) VALUES (\n2, 'Mirror', 20, 1, '2014-03-01T09:00:00Z',\n'{\"home-decor\", \"reflective\", \"glass\"}',\nNULL,\nNULL\n)",
+    "INSERT INTO \"myTestFlat\" (\n\"_id\", \"item\", \"price\", \"quantity\", \"date\", \"tags\", \"props\", \"sales\"\n) VALUES (\n3, 'Shampoo', 5, 10, '2014-03-15T09:00:00Z',\n'{\"hair-care\", \"personal-care\", \"premium\", \"herbal\"}',\n'{\"colors\": [\"Black\"], \"brand\": \"Sunsilk\", \"size\": \"L\", \"seller\": {\"name\": \"Metro Chemicals Pvt. Ltd.\", \"address\": {\"city\": \"Mumbai\", \"pincode\": 400004}}}',\nNULL\n)",
+    "INSERT INTO \"myTestFlat\" (\n\"_id\", \"item\", \"price\", \"quantity\", \"date\", \"tags\", \"props\", \"sales\"\n) VALUES (\n4, 'Shampoo', 5, 20, '2014-04-04T11:21:39.736Z',\n'{\"hair-care\", \"budget\", \"bulk\"}',\nNULL,\nNULL\n)",
+    "INSERT INTO \"myTestFlat\" (\n\"_id\", \"item\", \"price\", \"quantity\", \"date\", \"tags\", \"props\", \"sales\"\n) VALUES (\n5, 'Soap', 20, 5, '2014-04-04T21:23:13.331Z',\n'{\"hygiene\", \"antibacterial\", \"family-pack\"}',\n'{\"colors\": [\"Orange\", \"Blue\"], \"brand\": \"Lifebuoy\", \"size\": \"S\", \"seller\": {\"name\": \"Hans and Co.\", \"address\": {\"city\": \"Kolkata\", \"pincode\": 700007}}}',\nNULL\n)",
+    "INSERT INTO \"myTestFlat\" (\n\"_id\", \"item\", \"price\", \"quantity\", \"date\", \"tags\", \"props\", \"sales\"\n) VALUES (\n6, 'Comb', 7.5, 5, '2015-06-04T05:08:13Z',\n'{\"grooming\", \"plastic\", \"essential\"}',\nNULL,\nNULL\n)",
+    "INSERT INTO \"myTestFlat\" (\n\"_id\", \"item\", \"price\", \"quantity\", \"date\", \"tags\", \"props\", \"sales\"\n) VALUES (\n7, 'Comb', 7.5, 10, '2015-09-10T08:43:00Z',\n'{\"grooming\", \"bulk\", \"wholesale\"}',\n'{\"colors\": [], \"seller\": {\"name\": \"Go Go Plastics\", \"address\": {\"city\": \"Kolkata\", \"pincode\": 700007}}}',\nNULL\n)",
+    "INSERT INTO \"myTestFlat\" (\n\"_id\", \"item\", \"price\", \"quantity\", \"date\", \"tags\", \"props\", \"sales\"\n) VALUES (\n8, 'Soap', 10, 5, '2016-02-06T20:20:13Z',\n'{\"hygiene\", \"budget\", \"basic\"}',\nNULL,\nNULL\n)"
+  ]
+}

document-store/src/main/java/org/hypertrace/core/documentstore/Datastore.java

@@ -5,6 +5,7 @@
 import org.hypertrace.core.documentstore.metric.DocStoreMetricProvider;
 
 public interface Datastore {
+
   Set<String> listCollections();
 
   boolean createCollection(String collectionName, Map<String, String> options);
@@ -19,4 +20,19 @@ public interface Datastore {
   DocStoreMetricProvider getDocStoreMetricProvider();
 
   void close();
+
+  /**
+   * Returns a collection with the given name and type. A type can be used to specify different
+   * storage modes of the collection. For example, a collection can have all top-level fields or a
+   * single JSON column that contains all the fields. Both collections are handled differently in
+   * this case.
+   *
+   * @param collectionName name of the collection
+   * @param documentType type of the collection. For PG, we support FLAT and Legacy (for backward
+   *     compatibility)
+   * @return the corresponding collection impl
+   */
+  default Collection getCollectionForType(String collectionName, DocumentType documentType) {
+    throw new UnsupportedOperationException("Unsupported collection type: " + documentType);
+  }
 }

document-store/src/main/java/org/hypertrace/core/documentstore/mongo/MongoDatastore.java

@@ -15,6 +15,7 @@
 import org.bson.Document;
 import org.hypertrace.core.documentstore.Collection;
 import org.hypertrace.core.documentstore.Datastore;
+import org.hypertrace.core.documentstore.DocumentType;
 import org.hypertrace.core.documentstore.metric.DocStoreMetricProvider;
 import org.hypertrace.core.documentstore.metric.mongo.MongoDocStoreMetricProvider;
 import org.hypertrace.core.documentstore.model.config.ConnectionConfig;
@@ -25,6 +26,7 @@
 
 @Slf4j
 public class MongoDatastore implements Datastore {
+
   private static final Logger LOGGER = LoggerFactory.getLogger(MongoDatastore.class);
 
   private final ConnectionConfig connectionConfig;
@@ -78,6 +80,16 @@ public boolean deleteCollection(String collectionName) {
     return true;
   }
 
+  @Override
+  public Collection getCollectionForType(String collectionName, DocumentType documentType) {
+    // We support nested mongo docs
+    if (documentType == DocumentType.NESTED) {
+      return getCollection(collectionName);
+    }
+    throw new IllegalArgumentException(
+        "Unsupported document type: " + documentType + " for Mongo collection");
+  }
+
   @Override
   public Collection getCollection(String collectionName) {
     return new MongoCollection(

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

@@ -0,0 +1,57 @@
+package org.hypertrace.core.documentstore.postgres;
+
+import org.hypertrace.core.documentstore.CloseableIterator;
+import org.hypertrace.core.documentstore.Document;
+import org.hypertrace.core.documentstore.model.options.QueryOptions;
+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;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * PostgreSQL collection implementation for flat documents. All fields are stored as top-level
+ * PostgreSQL columns.
+ */
+// todo: Throw unsupported op exception for all write methods
+public class FlatPostgresCollection extends PostgresCollection {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(FlatPostgresCollection.class);
+
+  FlatPostgresCollection(final PostgresClient client, final String collectionName) {
+    super(client, collectionName);
+  }
+
+  @Override
+  public CloseableIterator<Document> query(
+      final org.hypertrace.core.documentstore.query.Query query, final QueryOptions queryOptions) {
+    PostgresQueryParser queryParser = createParser(query);
+    return queryWithParser(query, queryParser);
+  }
+
+  @Override
+  public CloseableIterator<Document> find(
+      final org.hypertrace.core.documentstore.query.Query query) {
+    PostgresQueryParser queryParser = createParser(query);
+    return queryWithParser(query, queryParser);
+  }
+
+  @Override
+  public long count(
+      org.hypertrace.core.documentstore.query.Query query, QueryOptions queryOptions) {
+    PostgresQueryParser queryParser =
+        new PostgresQueryParser(
+            tableIdentifier,
+            query,
+            new org.hypertrace.core.documentstore.postgres.query.v1.transformer
+                .FlatPostgresFieldTransformer());
+    return countWithParser(query, queryParser);
+  }
+
+  private PostgresQueryParser createParser(Query query) {
+    return new PostgresQueryParser(
+        tableIdentifier,
+        PostgresQueryExecutor.transformAndLog(query),
+        new FlatPostgresFieldTransformer());
+  }
+}

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

@@ -0,0 +1,51 @@
+package org.hypertrace.core.documentstore.postgres;
+
+import org.hypertrace.core.documentstore.CloseableIterator;
+import org.hypertrace.core.documentstore.Document;
+import org.hypertrace.core.documentstore.model.options.QueryOptions;
+import org.hypertrace.core.documentstore.postgres.query.v1.PostgresQueryParser;
+import org.hypertrace.core.documentstore.postgres.query.v1.transformer.NestedPostgresColTransformer;
+import org.hypertrace.core.documentstore.query.Query;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * PostgreSQL collection implementation for legacy document storage mode. Fields are stored within
+ * JSONB columns.
+ */
+public class NestedPostgresCollection extends PostgresCollection {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(NestedPostgresCollection.class);
+
+  NestedPostgresCollection(final PostgresClient client, final String collectionName) {
+    super(client, collectionName);
+  }
+
+  @Override
+  public CloseableIterator<Document> query(
+      final org.hypertrace.core.documentstore.query.Query query, final QueryOptions queryOptions) {
+    PostgresQueryParser queryParser = createParser(query);
+    return queryWithParser(query, queryParser);
+  }
+
+  @Override
+  public CloseableIterator<Document> find(
+      final org.hypertrace.core.documentstore.query.Query query) {
+    PostgresQueryParser queryParser = createParser(query);
+    return queryWithParser(query, queryParser);
+  }
+
+  @Override
+  public long count(
+      org.hypertrace.core.documentstore.query.Query query, QueryOptions queryOptions) {
+    PostgresQueryParser queryParser = createParser(query);
+    return countWithParser(query, queryParser);
+  }
+
+  private PostgresQueryParser createParser(Query query) {
+    return new PostgresQueryParser(
+        tableIdentifier,
+        PostgresQueryExecutor.transformAndLog(query),
+        new NestedPostgresColTransformer());
+  }
+}