SubDocUpdate Parsers for FlatPostgresCollection (#279)

Author: suddendust

document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java

@@ -1,9 +1,14 @@
 package org.hypertrace.core.documentstore.postgres;
 
+import static java.util.Map.entry;
 import static org.hypertrace.core.documentstore.model.options.ReturnDocumentType.AFTER_UPDATE;
 import static org.hypertrace.core.documentstore.model.options.ReturnDocumentType.BEFORE_UPDATE;
 import static org.hypertrace.core.documentstore.model.subdoc.UpdateOperator.ADD;
+import static org.hypertrace.core.documentstore.model.subdoc.UpdateOperator.ADD_TO_LIST_IF_ABSENT;
+import static org.hypertrace.core.documentstore.model.subdoc.UpdateOperator.APPEND_TO_LIST;
+import static org.hypertrace.core.documentstore.model.subdoc.UpdateOperator.REMOVE_ALL_FROM_LIST;
 import static org.hypertrace.core.documentstore.model.subdoc.UpdateOperator.SET;
+import static org.hypertrace.core.documentstore.model.subdoc.UpdateOperator.UNSET;
 
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
@@ -47,10 +52,14 @@
 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.postgres.query.v1.transformer.LegacyFilterToQueryFilterTransformer;
-import org.hypertrace.core.documentstore.postgres.update.FlatUpdateContext;
-import org.hypertrace.core.documentstore.postgres.update.parser.FlatCollectionSubDocAddOperatorParser;
-import org.hypertrace.core.documentstore.postgres.update.parser.FlatCollectionSubDocSetOperatorParser;
-import org.hypertrace.core.documentstore.postgres.update.parser.FlatCollectionSubDocUpdateOperatorParser;
+import org.hypertrace.core.documentstore.postgres.update.parser.PostgresAddToListIfAbsentParser;
+import org.hypertrace.core.documentstore.postgres.update.parser.PostgresAddValueParser;
+import org.hypertrace.core.documentstore.postgres.update.parser.PostgresAppendToListParser;
+import org.hypertrace.core.documentstore.postgres.update.parser.PostgresRemoveAllFromListParser;
+import org.hypertrace.core.documentstore.postgres.update.parser.PostgresSetValueParser;
+import org.hypertrace.core.documentstore.postgres.update.parser.PostgresUnsetPathParser;
+import org.hypertrace.core.documentstore.postgres.update.parser.PostgresUpdateOperationParser;
+import org.hypertrace.core.documentstore.postgres.update.parser.PostgresUpdateOperationParser.UpdateParserInput;
 import org.hypertrace.core.documentstore.postgres.utils.PostgresUtils;
 import org.hypertrace.core.documentstore.query.Query;
 import org.postgresql.util.PSQLException;
@@ -74,11 +83,14 @@ public class FlatPostgresCollection extends PostgresCollection {
   private static final String MISSING_COLUMN_STRATEGY_CONFIG = "missingColumnStrategy";
   private static final String DEFAULT_PRIMARY_KEY_COLUMN = "key";
 
-  private static final Map<UpdateOperator, FlatCollectionSubDocUpdateOperatorParser>
-      SUB_DOC_UPDATE_PARSERS =
-          Map.of(
-              SET, new FlatCollectionSubDocSetOperatorParser(),
-              ADD, new FlatCollectionSubDocAddOperatorParser());
+  private static final Map<UpdateOperator, PostgresUpdateOperationParser> UPDATE_PARSER_MAP =
+      Map.ofEntries(
+          entry(SET, new PostgresSetValueParser()),
+          entry(UNSET, new PostgresUnsetPathParser()),
+          entry(ADD, new PostgresAddValueParser()),
+          entry(REMOVE_ALL_FROM_LIST, new PostgresRemoveAllFromListParser()),
+          entry(ADD_TO_LIST_IF_ABSENT, new PostgresAddToListIfAbsentParser()),
+          entry(APPEND_TO_LIST, new PostgresAppendToListParser()));
 
   private final PostgresLazyilyLoadedSchemaRegistry schemaRegistry;
 
@@ -624,7 +636,7 @@ private Map<String, String> resolvePathsToColumns(
       UpdateOperator operator = update.getOperator();
 
       Preconditions.checkArgument(
-          SUB_DOC_UPDATE_PARSERS.containsKey(operator), "Unsupported UPDATE operator: " + operator);
+          UPDATE_PARSER_MAP.containsKey(operator), "Unsupported UPDATE operator: " + operator);
 
       String path = update.getSubDocument().getPath();
       Optional<String> columnName = resolveColumnName(path, tableName);
@@ -744,20 +756,40 @@ private void executeUpdate(
       PostgresColumnMetadata colMeta =
           schemaRegistry.getColumnOrRefresh(tableName, columnName).orElseThrow();
 
-      FlatUpdateContext context =
-          FlatUpdateContext.builder()
-              .columnName(columnName)
-              // get the nested path. So for example, if colName is `customAttr` and full path is
-              // `customAttr.props`, then the nested path is `props`.
-              .nestedPath(getNestedPath(path, columnName))
-              .columnType(colMeta.getPostgresType())
-              .value(update.getSubDocumentValue())
-              .params(params)
-              .build();
-
-      FlatCollectionSubDocUpdateOperatorParser operatorParser =
-          SUB_DOC_UPDATE_PARSERS.get(update.getOperator());
-      String fragment = operatorParser.parse(context);
+      String[] nestedPath = getNestedPath(path, columnName);
+      boolean isTopLevel = nestedPath.length == 0;
+      UpdateOperator operator = update.getOperator();
+
+      Params.Builder paramsBuilder = Params.newBuilder();
+      PostgresUpdateOperationParser unifiedParser = UPDATE_PARSER_MAP.get(operator);
+
+      String fragment;
+
+      if (isTopLevel) {
+        UpdateParserInput input =
+            UpdateParserInput.builder()
+                .baseField(columnName)
+                .path(new String[0])
+                .update(update)
+                .paramsBuilder(paramsBuilder)
+                .columnType(colMeta.getPostgresType())
+                .build();
+        fragment = unifiedParser.parseNonJsonbField(input);
+      } else {
+        // parseInternal() returns just the value expression
+        UpdateParserInput jsonbInput =
+            UpdateParserInput.builder()
+                .baseField(String.format("\"%s\"", columnName))
+                .path(nestedPath)
+                .update(update)
+                .paramsBuilder(paramsBuilder)
+                .columnType(colMeta.getPostgresType())
+                .build();
+        String valueExpr = unifiedParser.parseInternal(jsonbInput);
+        fragment = String.format("\"%s\" = %s", columnName, valueExpr);
+      }
+      // Transfer params from builder to our list
+      params.addAll(paramsBuilder.build().getObjectParams().values());
       setFragments.add(fragment);
     }
 

document-store/src/integrationTest/java/org/hypertrace/core/documentstore/MongoFlatPgConsistencyTest.java

@@ -6,6 +6,7 @@
 import org.hypertrace.core.documentstore.expression.impl.JsonIdentifierExpression;
 import org.hypertrace.core.documentstore.expression.impl.RelationalExpression;
 import org.hypertrace.core.documentstore.postgres.Params;
+import org.hypertrace.core.documentstore.postgres.query.v1.parser.filter.nonjson.field.PostgresDataType;
 
 /**
  * Optimized parser for IN operations on JSON array fields with type-specific casting.
@@ -119,7 +120,7 @@ private String prepareFilterStringForArrayInOperator(
       final Params.Builder paramsBuilder) {
 
     // Determine the appropriate type cast for jsonb_build_array elements
-    String typeCast = getTypeCastForArray(fieldType);
+    String typeCast = PostgresDataType.getJsonArrayElementTypeCast(fieldType);
 
     // For JSON arrays, we use the @> containment operator
     // Check if ANY of the RHS values is contained in the LHS array
@@ -137,26 +138,4 @@ private String prepareFilterStringForArrayInOperator(
         ? String.format("(%s)", orConditions)
         : orConditions;
   }
-
-  /**
-   * Returns the PostgreSQL type cast string for jsonb_build_array elements based on array type.
-   *
-   * @param fieldType The JSON field type (must not be null)
-   * @return Type cast string (e.g., "::text", "::numeric")
-   */
-  private String getTypeCastForArray(JsonFieldType fieldType) {
-    switch (fieldType) {
-      case STRING_ARRAY:
-        return "::text";
-      case NUMBER_ARRAY:
-        return "::numeric";
-      case BOOLEAN_ARRAY:
-        return "::boolean";
-      case OBJECT_ARRAY:
-        return "::jsonb";
-      default:
-        throw new IllegalArgumentException(
-            "Unsupported array type: " + fieldType + ". Expected *_ARRAY types.");
-    }
-  }
 }

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

@@ -1,6 +1,7 @@
 package org.hypertrace.core.documentstore.postgres.query.v1.parser.filter.nonjson.field;
 
 import org.hypertrace.core.documentstore.expression.impl.DataType;
+import org.hypertrace.core.documentstore.expression.impl.JsonFieldType;
 
 /**
  * PostgreSQL-specific data types with their SQL type strings.
@@ -40,6 +41,31 @@ public String getArraySqlType() {
     return sqlType + "[]";
   }
 
+  public String getTypeCast() {
+    return sqlType == null ? "" : "::" + sqlType;
+  }
+
+  public String getArrayTypeCast() {
+    return sqlType == null ? "" : "::" + sqlType + "[]";
+  }
+
+  public static PostgresDataType fromJavaValue(Object value) {
+    if (value instanceof String) {
+      return TEXT;
+    } else if (value instanceof Integer) {
+      return INTEGER;
+    } else if (value instanceof Long) {
+      return BIGINT;
+    } else if (value instanceof Float) {
+      return REAL;
+    } else if (value instanceof Double) {
+      return DOUBLE_PRECISION;
+    } else if (value instanceof Boolean) {
+      return BOOLEAN;
+    }
+    return UNKNOWN;
+  }
+
   /**
    * Maps a generic DataType to its PostgreSQL equivalent.
    *
@@ -70,4 +96,27 @@ public static PostgresDataType fromDataType(DataType dataType) {
         throw new IllegalArgumentException("Unknown DataType: " + dataType);
     }
   }
+
+  /**
+   * Returns the PostgreSQL type cast string for JSONB array element types.
+   *
+   * @param fieldType the JSON field type (must be an array type)
+   * @return Type cast string (e.g., "::text", "::numeric", "::boolean", "::jsonb")
+   * @throws IllegalArgumentException if fieldType is not a supported array type
+   */
+  public static String getJsonArrayElementTypeCast(JsonFieldType fieldType) {
+    switch (fieldType) {
+      case STRING_ARRAY:
+        return "::text";
+      case NUMBER_ARRAY:
+        return "::numeric";
+      case BOOLEAN_ARRAY:
+        return "::boolean";
+      case OBJECT_ARRAY:
+        return "::jsonb";
+      default:
+        throw new IllegalArgumentException(
+            "Unsupported array type: " + fieldType + ". Expected *_ARRAY types.");
+    }
+  }
 }

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

@@ -30,6 +30,7 @@
 import org.hypertrace.core.documentstore.postgres.query.v1.parser.builder.PostgresSelectExpressionParserBuilderImpl;
 import org.hypertrace.core.documentstore.postgres.query.v1.parser.filter.PostgresRelationalFilterParser.PostgresRelationalFilterContext;
 import org.hypertrace.core.documentstore.postgres.query.v1.parser.filter.PostgresRelationalFilterParserFactoryImpl;
+import org.hypertrace.core.documentstore.postgres.query.v1.parser.filter.nonjson.field.PostgresDataType;
 
 public class PostgresFilterTypeExpressionVisitor implements FilterTypeExpressionVisitor {
 
@@ -260,26 +261,20 @@ private String inferArrayTypeCastFromFilter(FilterTypeExpression filter) {
     if (filter instanceof RelationalExpression) {
       RelationalExpression relExpr = (RelationalExpression) filter;
 
-      // The visitor returns a string representation, but we need the actual value
       // Try to get the constant value directly if it's a ConstantExpression
       if (relExpr.getRhs() instanceof ConstantExpression) {
         ConstantExpression constExpr = (ConstantExpression) relExpr.getRhs();
         Object value = constExpr.getValue();
 
-        if (value instanceof String) {
-          return "::text[]";
-        } else if (value instanceof Integer || value instanceof Long) {
-          return "::bigint[]";
-        } else if (value instanceof Double || value instanceof Float) {
-          return "::double precision[]";
-        } else if (value instanceof Boolean) {
-          return "::boolean[]";
+        PostgresDataType pgType = PostgresDataType.fromJavaValue(value);
+        if (pgType != PostgresDataType.UNKNOWN) {
+          return pgType.getArrayTypeCast();
         }
       }
     }
 
     // Default to text[] if we can't infer the type
-    return "::text[]";
+    return PostgresDataType.TEXT.getArrayTypeCast();
   }
 
   private String getFilterStringForAnyOperator(final DocumentArrayFilterExpression expression) {