FlatPostgresCollection createOrReplace #268

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open

suddendust wants to merge 42 commits into hypertrace:main from suddendust:pg_write_createOrReplace

...e/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java

Large diffs are not rendered by default.

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

-Original file line number
+Diff line change
@@ Expand Up / @@ -23,4 +23,6 @@ public interface ColumnMetadata { @@
        * @return whether this column is an array type
        */
       boolean isArray();
+      boolean isPrimaryKey();
     }

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

-Original file line number
+Diff line change
@@ Expand Up / @@ -38,4 +38,12 @@ public interface SchemaRegistry<T extends ColumnMetadata> { @@
        * @return optional of the col metadata.
        */
       Optional<T> getColumnOrRefresh(String tableName, String colName);
+      /**
+       * Returns the primary key column name for the given table.
+       *
+       * @param tableName the table name
+       * @return optional of the primary key column name
+       */
+      Optional<String> getPrimaryKeyColumn(String tableName);
     }

.../src/main/java/org/hypertrace/core/documentstore/model/options/MissingColumnStrategy.java

-Original file line number
+Diff line change
@@ Expand Up / @@ -19,5 +19,10 @@ public enum MissingColumnStrategy { @@
        * a field doesn't match the schema. The write operation will fail.
        */
       THROW,
-      IGNORE_DOCUMENT
+      /** Ignore the entire document if it doesn't match the schema. */
+      IGNORE_DOCUMENT;
+      public static MissingColumnStrategy defaultStrategy() {
+        return SKIP;
+      }
     }

...tore/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java

            
                      Original file line number
                      Diff line number
                      Diff line change
                  
    @@ -8,6 +8,7 @@
  
    import com.fasterxml.jackson.databind.ObjectMapper;

    import com.google.common.base.Preconditions;

    import java.io.IOException;

    import java.sql.Array;

    import java.sql.Connection;

    import java.sql.Date;

    import java.sql.PreparedStatement;

    @@ -16,15 +17,9 @@
  
    import java.sql.Timestamp;

    import java.sql.Types;

    import java.time.Instant;

    import java.util.ArrayList;

    import java.util.Collection;

    import java.util.HashMap;

    import java.util.Iterator;

    import java.util.List;

    import java.util.Map;

    import java.util.*;

    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;

    @@ -52,6 +47,7 @@
  
    import org.hypertrace.core.documentstore.postgres.update.FlatUpdateContext;

    import org.hypertrace.core.documentstore.postgres.update.parser.FlatCollectionSubDocSetOperatorParser;

    import org.hypertrace.core.documentstore.postgres.update.parser.FlatCollectionSubDocUpdateOperatorParser;

    import org.hypertrace.core.documentstore.postgres.utils.PostgresUtils;

    import org.hypertrace.core.documentstore.query.Query;

    import org.postgresql.util.PSQLException;

    import org.postgresql.util.PSQLState;

    @@ -72,6 +68,7 @@ public class FlatPostgresCollection extends PostgresCollection {
  
      private static final String WRITE_NOT_SUPPORTED =

          "Write operations are not supported for flat collections yet!";

      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());

    @@ -88,24 +85,35 @@ public class FlatPostgresCollection extends PostgresCollection {
  
          final PostgresClient client,

          final String collectionName,

          final PostgresLazyilyLoadedSchemaRegistry schemaRegistry) {

        this(client, collectionName, schemaRegistry, null);

      }

      FlatPostgresCollection(

          final PostgresClient client,

          final String collectionName,

          final PostgresLazyilyLoadedSchemaRegistry schemaRegistry,

          final MissingColumnStrategy missingColumnStrategy) {

        super(client, collectionName);

        this.schemaRegistry = schemaRegistry;

        this.missingColumnStrategy = parseMissingColumnStrategy(client.getCustomParameters());

        this.missingColumnStrategy =

            missingColumnStrategy != null

                ? missingColumnStrategy

                : parseMissingColumnStrategy(client.getCustomParameters());

      }

      private static MissingColumnStrategy parseMissingColumnStrategy(Map<String, String> params) {

        String value = params.get(MISSING_COLUMN_STRATEGY_CONFIG);

        if (value == null || value.isEmpty()) {

          return MissingColumnStrategy.SKIP; // default

          return MissingColumnStrategy.defaultStrategy();

        }

        try {

          return MissingColumnStrategy.valueOf(value.toUpperCase());

        } catch (IllegalArgumentException e) {

          LOGGER.warn(

              "Invalid missingColumnStrategy value: '{}', using default SKIP. Valid values: {}",

              value,

              java.util.Arrays.toString(MissingColumnStrategy.values()));

          return MissingColumnStrategy.SKIP;

              Arrays.toString(MissingColumnStrategy.values()));

          return MissingColumnStrategy.defaultStrategy();

        }

      }

    @@ -215,7 +223,7 @@ public CreateResult create(Key key, Document document) throws IOException {
  
      @Override

      public boolean createOrReplace(Key key, Document document) throws IOException {

        throw new UnsupportedOperationException(WRITE_NOT_SUPPORTED);

        return createOrReplaceWithRetry(key, document, false);

      }

      @Override

    @@ -487,6 +495,12 @@ private CreateResult createWithRetry(Key key, Document document, boolean isRetry
  
        try {

          TypedDocument parsed = parseDocument(document, tableName, skippedFields);

          // Add the key as the primary key column

          String pkColumn = getPKForTable(tableName);

          String quotedPkColumn = PostgresUtils.wrapFieldNamesWithDoubleQuotes(pkColumn);

          PostgresDataType pkType = getPrimaryKeyType(tableName, pkColumn);

          parsed.add(quotedPkColumn, key.toString(), pkType, false);

          // If IGNORE_DOCUMENT strategy and any fields were skipped, ignore the entire document

          if (missingColumnStrategy == MissingColumnStrategy.IGNORE_DOCUMENT

              && !skippedFields.isEmpty()) {

    @@ -495,12 +509,6 @@ private CreateResult createWithRetry(Key key, Document document, boolean isRetry
  
            return new CreateResult(CreateStatus.IGNORED, isRetry, 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(CreateStatus.FAILED, isRetry, skippedFields);

          }

          String sql = buildInsertSql(parsed.getColumns());

          LOGGER.debug("Insert SQL: {}", sql);

    @@ -543,12 +551,18 @@ private TypedDocument parseDocument(
  
            continue;

          }

          if (columnMetadata.get().isPrimaryKey()) {

            // PK is added by the caller

            continue;

          }

          PostgresDataType type = columnMetadata.get().getPostgresType();

          boolean isArray = columnMetadata.get().isArray();

          try {

            Object value = extractValue(fieldValue, type, isArray);

            typedDocument.add("\"" + fieldName + "\"", value, type, isArray);

            typedDocument.add(

                PostgresUtils.wrapFieldNamesWithDoubleQuotes(fieldName), value, type, isArray);

          } catch (Exception e) {

            if (missingColumnStrategy == MissingColumnStrategy.THROW) {

              throw new SchemaMismatchException(

    @@ -589,12 +603,132 @@ private int executeUpdate(String sql, TypedDocument parsed) throws SQLException
  
        }

      }

      private boolean createOrReplaceWithRetry(Key key, Document document, boolean isRetry)

          throws IOException {

        String tableName = tableIdentifier.getTableName();

        List<String> skippedFields = new ArrayList<>();

        try {

          TypedDocument parsed = parseDocument(document, tableName, skippedFields);

          // Add the key as the primary key column

          String pkColumn = getPKForTable(tableName);

          String quotedPkColumn = PostgresUtils.wrapFieldNamesWithDoubleQuotes(pkColumn);

          PostgresDataType pkType = getPrimaryKeyType(tableName, pkColumn);

          parsed.add(quotedPkColumn, key.toString(), pkType, false);

          String sql = buildUpsertSql(parsed.getColumns(), quotedPkColumn);

          LOGGER.debug("Upsert SQL: {}", sql);

          return executeUpsert(sql, parsed);

        } catch (PSQLException e) {

          return handlePSQLExceptionForUpsert(e, key, document, tableName, isRetry);

        } catch (SQLException e) {

          LOGGER.error("SQLException in createOrReplace. key: {} content: {}", key, document, e);

          throw new IOException(e);

        }

      }

      /**

       * Builds a PostgreSQL upsert (INSERT ... ON CONFLICT DO UPDATE) SQL statement.

       *

       * <p>This method constructs an atomic upsert query that:

       *

       * <ul>

       *   <li>Inserts a new row if no conflict on the primary key

       *   <li>Updates all non-PK columns if a row with the same PK already exists

       * </ul>

       *

       * <p><b>Generated SQL pattern:</b>

       *

       * <pre>{@code

       * INSERT INTO table (col1, col2, pk_col)

       * VALUES (?, ?, ?)

       * ON CONFLICT (pk_col) DO UPDATE SET col1 = EXCLUDED.col1, col2 = EXCLUDED.col2

       * RETURNING (xmax = 0) AS is_insert

       * }</pre>

       *

       * <p><b>The EXCLUDED table:</b> In PostgreSQL's ON CONFLICT clause, {@code EXCLUDED} is a special

       * table that references the row that would have been inserted (the "proposed" row). This allows

       * us to update existing rows with the new values without re-specifying them.

       *

       * <p><b>The RETURNING clause:</b> {@code (xmax = 0) AS is_insert} is a PostgreSQL trick to

       * determine if the operation was an INSERT or UPDATE:

       *

       * <ul>

       *   <li>{@code xmax} is a system column that stores the transaction ID of the deleting/updating

       *       transaction

       *   <li>For a freshly inserted row, {@code xmax = 0} (no prior transaction modified it)

       *   <li>For an updated row, {@code xmax != 0} (the UPDATE sets it to the current transaction ID)

       *   <li>Thus, {@code is_insert = true} means INSERT, {@code is_insert = false} means UPDATE

       * </ul>

       *

       * @param columns List of quoted column names to include in the upsert (including PK)

       * @param pkColumn The quoted primary key column name used for conflict detection

       * @return The complete upsert SQL statement with placeholders for values

       */

      private String buildUpsertSql(List<String> columns, String pkColumn) {

        String columnList = String.join(", ", columns);

        String placeholders = String.join(", ", columns.stream().map(c -> "?").toArray(String[]::new));

        // Build SET clause for non-PK columns: col = EXCLUDED.col

        String setClause =

            columns.stream()

                .filter(col -> !col.equals(pkColumn))

                .map(col -> col + " = EXCLUDED." + col)

                .collect(Collectors.joining(", "));

        return String.format(

            "INSERT INTO %s (%s) VALUES (%s) ON CONFLICT (%s) DO UPDATE SET %s RETURNING (xmax = 0) AS is_insert",

            tableIdentifier, columnList, placeholders, pkColumn, setClause);

      }

      private boolean executeUpsert(String sql, TypedDocument parsed) throws SQLException {

        try (Connection conn = client.getPooledConnection();

            PreparedStatement ps = conn.prepareStatement(sql)) {

          int index = 1;

          for (String column : parsed.getColumns()) {

            setParameter(

                conn,

                ps,

                index++,

                parsed.getValue(column),

                parsed.getType(column),

                parsed.isArray(column));

          }

          try (ResultSet rs = ps.executeQuery()) {

            if (rs.next()) {

              // is_insert is true if xmax = 0 (new row), false if updated. This helps us differentiate

              // b/w creates/upserts

              return rs.getBoolean("is_insert");

            }

          }

          return false;

        }

      }

      private boolean handlePSQLExceptionForUpsert(

          PSQLException e, Key key, Document document, String tableName, boolean isRetry)

          throws IOException {

        if (!isRetry && shouldRefreshSchemaAndRetry(e.getSQLState())) {

          LOGGER.info(

              "Schema mismatch detected during upsert (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 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: {}",

              "Schema mismatch detected during create (SQLState: {}), refreshing schema and retrying. key: {}",

              e.getSQLState(),

              key);

          schemaRegistry.invalidate(tableName);

    @@ -613,6 +747,17 @@ private boolean shouldRefreshSchemaAndRetry(String sqlState) {
  
            || PSQLState.DATATYPE_MISMATCH.getState().equals(sqlState);

      }

      private String getPKForTable(String tableName) {

        return schemaRegistry.getPrimaryKeyColumn(tableName).orElse(DEFAULT_PRIMARY_KEY_COLUMN);

      }

      private PostgresDataType getPrimaryKeyType(String tableName, String pkColumn) {

        return schemaRegistry

            .getColumnOrRefresh(tableName, pkColumn)

            .map(PostgresColumnMetadata::getPostgresType)

            .orElse(PostgresDataType.TEXT);

      }

      /**

       * Typed document contains field information along with the field type. Uses LinkedHashMaps keyed

       * by column name. LinkedHashMap preserves insertion order for consistent parameter binding.

    @@ -629,10 +774,6 @@ void add(String column, Object value, PostgresDataType type, boolean isArray) {
  
          arrays.put(column, isArray);

        }

        boolean isEmpty() {

          return values.isEmpty();

        }

        List<String> getColumns() {

          return new ArrayList<>(values.keySet());

        }

    @@ -721,8 +862,9 @@ private void setParameter(
  
        }

        if (isArray) {

          Object[] arrayValues = (value instanceof Object[]) ? (Object[]) value : new Object[] {value};

          java.sql.Array sqlArray = conn.createArrayOf(type.getSqlType(), arrayValues);

          // todo: Maybe check if the value is actually an array

          Object[] arrayValues = (Object[]) value;

          Array sqlArray = conn.createArrayOf(type.getSqlType(), arrayValues);

          ps.setArray(index, sqlArray);

          return;

        }

...ent-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresDatastore.java

-Original file line number
+Diff line change
@@ Expand Up / @@ -26,6 +26,7 @@ @@
     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.model.options.MissingColumnStrategy;
     import org.hypertrace.core.documentstore.postgres.model.PostgresColumnMetadata;
     import org.slf4j.Logger;
     import org.slf4j.LoggerFactory;
@@ Expand Down Expand Up @@
         }
       }
+      public Collection getFlatCollection(
+          String collectionName, MissingColumnStrategy missingColumnStrategy) {
+        return new FlatPostgresCollection(
+            client,
+            collectionName,
+            (PostgresLazyilyLoadedSchemaRegistry) schemaRegistry,
+            missingColumnStrategy);
+      }
       @Override
       public boolean healthCheck() {
         String healthCheckSql = "SELECT 1;";
@@ Expand Down @@

.../java/org/hypertrace/core/documentstore/postgres/PostgresLazyilyLoadedSchemaRegistry.java

-Original file line number
+Diff line change
@@ Expand Up / @@ -170,4 +170,19 @@ private boolean canRefresh(String tableName) { @@
         }
         return Duration.between(lastRefresh, Instant.now()).compareTo(refreshCooldown) >= 0;
       }
+      /**
+       * Returns the primary key column name for the given table.
+       *
+       * @param tableName the name of the table
+       * @return optional of the primary key column name, or empty if no primary key is found
+       */
+      @Override
+      public Optional<String> getPrimaryKeyColumn(String tableName) {
+        Map<String, PostgresColumnMetadata> schema = getSchema(tableName);
+        return schema.values().stream()
+            .filter(PostgresColumnMetadata::isPrimaryKey)
+            .map(PostgresColumnMetadata::getName)
+            .findFirst();
+      }
     }

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

FlatPostgresCollection createOrReplace #268

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!

Uh oh!

Uh oh!

FlatPostgresCollection createOrReplace #268

Are you sure you want to change the base?

Uh oh!

FlatPostgresCollection createOrReplace #268

Uh oh!

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!

Uh oh!

Uh oh!