envio 2.26.0-alpha.9 → 2.26.0-rc.0

package/src/PgStorage.res

@@ -1,16 +1,16 @@
-let makeCreateIndexQuery = (~tableName, ~indexFields, ~pgSchema) => {
+let makeCreateIndexSql = (~tableName, ~indexFields, ~pgSchema) => {
   let indexName = tableName ++ "_" ++ indexFields->Js.Array2.joinWith("_")
   let index = indexFields->Belt.Array.map(idx => `"${idx}"`)->Js.Array2.joinWith(", ")
   `CREATE INDEX IF NOT EXISTS "${indexName}" ON "${pgSchema}"."${tableName}"(${index});`
 }
 
-let makeCreateTableIndicesQuery = (table: Table.table, ~pgSchema) => {
+let makeCreateTableIndicesSql = (table: Table.table, ~pgSchema) => {
   open Belt
   let tableName = table.tableName
   let createIndex = indexField =>
-    makeCreateIndexQuery(~tableName, ~indexFields=[indexField], ~pgSchema)
+    makeCreateIndexSql(~tableName, ~indexFields=[indexField], ~pgSchema)
   let createCompositeIndex = indexFields => {
-    makeCreateIndexQuery(~tableName, ~indexFields, ~pgSchema)
+    makeCreateIndexSql(~tableName, ~indexFields, ~pgSchema)
   }
 
   let singleIndices = table->Table.getSingleIndices
@@ -20,7 +20,7 @@ let makeCreateTableIndicesQuery = (table: Table.table, ~pgSchema) => {
     compositeIndices->Array.map(createCompositeIndex)->Js.Array2.joinWith("\n")
 }
 
-let makeCreateTableQuery = (table: Table.table, ~pgSchema) => {
+let makeCreateTableSql = (table: Table.table, ~pgSchema) => {
   open Belt
   let fieldsMapped =
     table
@@ -58,7 +58,7 @@ let makeInitializeTransaction = (
   ~generalTables=[],
   ~entities=[],
   ~enums=[],
-  ~reuseExistingPgSchema=false,
+  ~cleanRun=false,
 ) => {
   let allTables = generalTables->Array.copy
   let allEntityTables = []
@@ -71,15 +71,12 @@ let makeInitializeTransaction = (
 
   let query = ref(
     (
-      reuseExistingPgSchema
-      // Hosted Service already have a DB with the created public schema
-      // It also doesn't allow to simply drop it,
-      // so we reuse an existing schema when it's empty (our case)
-        ? ""
-        : `DROP SCHEMA IF EXISTS "${pgSchema}" CASCADE;
-CREATE SCHEMA "${pgSchema}";\n`
+      cleanRun
+        ? `DROP SCHEMA IF EXISTS "${pgSchema}" CASCADE;
+CREATE SCHEMA "${pgSchema}";`
+        : `CREATE SCHEMA IF NOT EXISTS "${pgSchema}";`
     ) ++
-    `GRANT ALL ON SCHEMA "${pgSchema}" TO "${pgUser}";
+    `GRANT ALL ON SCHEMA "${pgSchema}" TO ${pgUser};
 GRANT ALL ON SCHEMA "${pgSchema}" TO public;`,
   )
 
@@ -90,17 +87,31 @@ GRANT ALL ON SCHEMA "${pgSchema}" TO public;`,
       ->Js.Array2.map(v => `'${v->(Utils.magic: Internal.enum => string)}'`)
       ->Js.Array2.joinWith(", ")});`
 
-    query := query.contents ++ "\n" ++ enumCreateQuery
+    query :=
+      query.contents ++
+      "\n" ++ if cleanRun {
+        // Direct creation when cleanRunting (faster)
+        enumCreateQuery
+      } else {
+        // Wrap with conditional check only when not cleanRunting
+        `IF NOT EXISTS (
+  SELECT 1 FROM pg_type 
+  WHERE typname = '${enumConfig.name->Js.String2.toLowerCase}' 
+  AND typnamespace = (SELECT oid FROM pg_namespace WHERE nspname = '${pgSchema}')
+) THEN 
+    ${enumCreateQuery}
+END IF;`
+      }
   })
 
   // Batch all table creation first (optimal for PostgreSQL)
   allTables->Js.Array2.forEach((table: Table.table) => {
-    query := query.contents ++ "\n" ++ makeCreateTableQuery(table, ~pgSchema)
+    query := query.contents ++ "\n" ++ makeCreateTableSql(table, ~pgSchema)
   })
 
   // Then batch all indices (better performance when tables exist)
   allTables->Js.Array2.forEach((table: Table.table) => {
-    let indices = makeCreateTableIndicesQuery(table, ~pgSchema)
+    let indices = makeCreateTableIndicesSql(table, ~pgSchema)
     if indices !== "" {
       query := query.contents ++ "\n" ++ indices
     }
@@ -120,7 +131,7 @@ GRANT ALL ON SCHEMA "${pgSchema}" TO public;`,
       query :=
         query.contents ++
         "\n" ++
-        makeCreateIndexQuery(
+        makeCreateIndexSql(
           ~tableName=derivedFromField.derivedFromEntity,
           ~indexFields=[indexField],
           ~pgSchema,
@@ -128,26 +139,28 @@ GRANT ALL ON SCHEMA "${pgSchema}" TO public;`,
     })
   })
 
-  [query.contents]->Js.Array2.concat(
-    functionsQuery.contents !== "" ? [functionsQuery.contents] : [],
-  )
+  [
+    // Return optimized queries - main DDL in DO block, functions separate
+    // Note: DO $$ BEGIN wrapper is only needed for PL/pgSQL conditionals (IF NOT EXISTS)
+    // Reset case uses direct DDL (faster), non-cleanRun case uses conditionals (safer)
+    cleanRun || enums->Utils.Array.isEmpty
+      ? query.contents
+      : `DO $$ BEGIN ${query.contents} END $$;`,
+    // Functions query (separate as they can't be in DO block)
+  ]->Js.Array2.concat(functionsQuery.contents !== "" ? [functionsQuery.contents] : [])
 }
 
-let makeLoadByIdQuery = (~pgSchema, ~tableName) => {
+let makeLoadByIdSql = (~pgSchema, ~tableName) => {
   `SELECT * FROM "${pgSchema}"."${tableName}" WHERE id = $1 LIMIT 1;`
 }
 
-let makeLoadByFieldQuery = (~pgSchema, ~tableName, ~fieldName, ~operator) => {
-  `SELECT * FROM "${pgSchema}"."${tableName}" WHERE "${fieldName}" ${operator} $1;`
-}
-
-let makeLoadByIdsQuery = (~pgSchema, ~tableName) => {
+let makeLoadByIdsSql = (~pgSchema, ~tableName) => {
   `SELECT * FROM "${pgSchema}"."${tableName}" WHERE id = ANY($1::text[]);`
 }
 
-let makeInsertUnnestSetQuery = (~pgSchema, ~table: Table.table, ~itemSchema, ~isRawEvents) => {
+let makeInsertUnnestSetSql = (~pgSchema, ~table: Table.table, ~itemSchema, ~isRawEvents) => {
   let {quotedFieldNames, quotedNonPrimaryFieldNames, arrayFieldTypes} =
-    table->Table.toSqlParams(~schema=itemSchema, ~pgSchema)
+    table->Table.toSqlParams(~schema=itemSchema)
 
   let primaryKeyFieldNames = Table.getPrimaryKeyFieldNames(table)
 
@@ -175,9 +188,8 @@ SELECT * FROM unnest(${arrayFieldTypes
   } ++ ";"
 }
 
-let makeInsertValuesSetQuery = (~pgSchema, ~table: Table.table, ~itemSchema, ~itemsCount) => {
-  let {quotedFieldNames, quotedNonPrimaryFieldNames} =
-    table->Table.toSqlParams(~schema=itemSchema, ~pgSchema)
+let makeInsertValuesSetSql = (~pgSchema, ~table: Table.table, ~itemSchema, ~itemsCount) => {
+  let {quotedFieldNames, quotedNonPrimaryFieldNames} = table->Table.toSqlParams(~schema=itemSchema)
 
   let primaryKeyFieldNames = Table.getPrimaryKeyFieldNames(table)
   let fieldsCount = quotedFieldNames->Array.length
@@ -224,13 +236,12 @@ VALUES${placeholders.contents}` ++
 // they are always guaranteed to be an object.
 // FIXME what about Fuel params?
 let rawEventsTableName = "raw_events"
-let eventSyncStateTableName = "event_sync_state"
 
 // Constants for chunking
 let maxItemsPerQuery = 500
 
 let makeTableBatchSetQuery = (~pgSchema, ~table: Table.table, ~itemSchema: S.t<'item>) => {
-  let {dbSchema, hasArrayField} = table->Table.toSqlParams(~schema=itemSchema, ~pgSchema)
+  let {dbSchema, hasArrayField} = table->Table.toSqlParams(~schema=itemSchema)
   let isRawEvents = table.tableName === rawEventsTableName
 
   // Should experiment how much it'll affect performance
@@ -242,7 +253,7 @@ let makeTableBatchSetQuery = (~pgSchema, ~table: Table.table, ~itemSchema: S.t<'
 
   if isRawEvents || !hasArrayField {
     {
-      "query": makeInsertUnnestSetQuery(~pgSchema, ~table, ~itemSchema, ~isRawEvents),
+      "sql": makeInsertUnnestSetSql(~pgSchema, ~table, ~itemSchema, ~isRawEvents),
       "convertOrThrow": S.compile(
         S.unnest(dbSchema),
         ~input=Value,
@@ -254,12 +265,7 @@ let makeTableBatchSetQuery = (~pgSchema, ~table: Table.table, ~itemSchema: S.t<'
     }
   } else {
     {
-      "query": makeInsertValuesSetQuery(
-        ~pgSchema,
-        ~table,
-        ~itemSchema,
-        ~itemsCount=maxItemsPerQuery,
-      ),
+      "sql": makeInsertValuesSetSql(~pgSchema, ~table, ~itemSchema, ~itemsCount=maxItemsPerQuery),
       "convertOrThrow": S.compile(
         S.unnest(itemSchema)->S.preprocess(_ => {
           serializer: Utils.Array.flatten->Utils.magic,
@@ -285,35 +291,6 @@ let chunkArray = (arr: array<'a>, ~chunkSize) => {
   chunks
 }
 
-let removeInvalidUtf8InPlace = entities =>
-  entities->Js.Array2.forEach(item => {
-    let dict = item->(Utils.magic: 'a => dict<unknown>)
-    dict->Utils.Dict.forEachWithKey((key, value) => {
-      if value->Js.typeof === "string" {
-        let value = value->(Utils.magic: unknown => string)
-        // We mutate here, since we don't care
-        // about the original value with \x00 anyways.
-        //
-        // This is unsafe, but we rely that it'll use
-        // the mutated reference on retry.
-        // TODO: Test it properly after we start using
-        // in-memory PGLite for indexer test framework.
-        dict->Js.Dict.set(
-          key,
-          value
-          ->Utils.String.replaceAll("\x00", "")
-          ->(Utils.magic: string => unknown),
-        )
-      }
-    })
-  })
-
-let pgEncodingErrorSchema = S.object(s =>
-  s.tag("message", `invalid byte sequence for encoding "UTF8": 0x00`)
-)
-
-exception PgEncodingError({table: Table.table})
-
 // WeakMap for caching table batch set queries
 let setQueryCache = Utils.WeakMap.make()
 let setOrThrow = async (sql, ~items, ~table: Table.table, ~itemSchema, ~pgSchema) => {
@@ -321,7 +298,7 @@ let setOrThrow = async (sql, ~items, ~table: Table.table, ~itemSchema, ~pgSchema
     ()
   } else {
     // Get or create cached query for this table
-    let data = switch setQueryCache->Utils.WeakMap.get(table) {
+    let query = switch setQueryCache->Utils.WeakMap.get(table) {
     | Some(cached) => cached
     | None => {
         let newQuery = makeTableBatchSetQuery(
@@ -334,31 +311,50 @@ let setOrThrow = async (sql, ~items, ~table: Table.table, ~itemSchema, ~pgSchema
       }
     }
 
+    let sqlQuery = query["sql"]
+
     try {
-      if data["isInsertValues"] {
-        let chunks = chunkArray(items, ~chunkSize=maxItemsPerQuery)
+      let payload =
+        query["convertOrThrow"](items->(Utils.magic: array<'item> => array<unknown>))->(
+          Utils.magic: unknown => array<unknown>
+        )
+
+      if query["isInsertValues"] {
+        let fieldsCount = switch itemSchema->S.classify {
+        | S.Object({items}) => items->Array.length
+        | _ => Js.Exn.raiseError("Expected an object schema for table")
+        }
+
+        // Chunk the items for VALUES-based queries
+        // We need to multiply by fields count,
+        // because we flattened our entity values with S.unnest
+        // to optimize the query execution.
+        let maxChunkSize = maxItemsPerQuery * fieldsCount
+        let chunks = chunkArray(payload, ~chunkSize=maxChunkSize)
         let responses = []
         chunks->Js.Array2.forEach(chunk => {
           let chunkSize = chunk->Array.length
-          let isFullChunk = chunkSize === maxItemsPerQuery
+          let isFullChunk = chunkSize === maxChunkSize
 
           let response = sql->Postgres.preparedUnsafe(
             // Either use the sql query for full chunks from cache
             // or create a new one for partial chunks on the fly.
             isFullChunk
-              ? data["query"]
-              : makeInsertValuesSetQuery(~pgSchema, ~table, ~itemSchema, ~itemsCount=chunkSize),
-            data["convertOrThrow"](chunk->(Utils.magic: array<'item> => array<unknown>)),
+              ? sqlQuery
+              : makeInsertValuesSetSql(
+                  ~pgSchema,
+                  ~table,
+                  ~itemSchema,
+                  ~itemsCount=chunkSize / fieldsCount,
+                ),
+            chunk->Utils.magic,
           )
           responses->Js.Array2.push(response)->ignore
         })
         let _ = await Promise.all(responses)
       } else {
         // Use UNNEST approach for single query
-        await sql->Postgres.preparedUnsafe(
-          data["query"],
-          data["convertOrThrow"](items->(Utils.magic: array<'item> => array<unknown>)),
-        )
+        await sql->Postgres.preparedUnsafe(sqlQuery, payload->Obj.magic)
       }
     } catch {
     | S.Raised(_) as exn =>
@@ -379,81 +375,23 @@ let setOrThrow = async (sql, ~items, ~table: Table.table, ~itemSchema, ~pgSchema
   }
 }
 
-let setEntityHistoryOrThrow = (
-  sql,
-  ~entityHistory: EntityHistory.t<'entity>,
-  ~rows: array<EntityHistory.historyRow<'entity>>,
-  ~shouldCopyCurrentEntity=?,
-  ~shouldRemoveInvalidUtf8=false,
-) => {
-  rows
-  ->Belt.Array.map(historyRow => {
-    let row = historyRow->S.reverseConvertToJsonOrThrow(entityHistory.schema)
-    if shouldRemoveInvalidUtf8 {
-      [row]->removeInvalidUtf8InPlace
-    }
-    entityHistory.insertFn(
-      sql,
-      row,
-      ~shouldCopyCurrentEntity=switch shouldCopyCurrentEntity {
-      | Some(v) => v
-      | None => {
-          let containsRollbackDiffChange =
-            historyRow.containsRollbackDiffChange->Belt.Option.getWithDefault(false)
-          !containsRollbackDiffChange
-        }
-      },
-    )
-  })
-  ->Promise.all
-  ->(Utils.magic: promise<array<unit>> => promise<unit>)
-}
-
-type schemaTableName = {
-  @as("table_name")
-  tableName: string,
-}
-
-let makeSchemaTableNamesQuery = (~pgSchema) => {
-  `SELECT table_name FROM information_schema.tables WHERE table_schema = '${pgSchema}';`
-}
-
 let make = (~sql: Postgres.sql, ~pgSchema, ~pgUser): Persistence.storage => {
   let isInitialized = async () => {
-    let envioTables =
+    let schemas =
       await sql->Postgres.unsafe(
-        `SELECT table_schema FROM information_schema.tables WHERE table_schema = '${pgSchema}' AND table_name = '${eventSyncStateTableName}';`,
+        `SELECT schema_name FROM information_schema.schemata WHERE schema_name = '${pgSchema}';`,
       )
-    envioTables->Utils.Array.notEmpty
+    schemas->Utils.Array.notEmpty
   }
 
-  let initialize = async (~entities=[], ~generalTables=[], ~enums=[]) => {
-    let schemaTableNames: array<schemaTableName> =
-      await sql->Postgres.unsafe(makeSchemaTableNamesQuery(~pgSchema))
-
-    // The initialization query will completely drop the schema and recreate it from scratch.
-    // So we need to check if the schema is not used for anything else than envio.
-    if (
-      // Should pass with existing schema with no tables
-      // This might happen when used with public schema
-      // which is automatically created by postgres.
-      schemaTableNames->Utils.Array.notEmpty &&
-        // Otherwise should throw if there's a table, but no envio specific one
-        // This means that the schema is used for something else than envio.
-        !(schemaTableNames->Js.Array2.some(table => table.tableName === eventSyncStateTableName))
-    ) {
-      Js.Exn.raiseError(
-        `Cannot run Envio migrations on PostgreSQL schema "${pgSchema}" because it contains non-Envio tables. Running migrations would delete all data in this schema.\n\nTo resolve this:\n1. If you want to use this schema, first backup any important data, then drop it with: "pnpm envio local db-migrate down"\n2. Or specify a different schema name by setting the "ENVIO_PG_PUBLIC_SCHEMA" environment variable\n3. Or manually drop the schema in your database if you're certain the data is not needed.`,
-      )
-    }
-
+  let initialize = async (~entities=[], ~generalTables=[], ~enums=[], ~cleanRun=false) => {
     let queries = makeInitializeTransaction(
       ~pgSchema,
       ~pgUser,
       ~generalTables,
       ~entities,
       ~enums,
-      ~reuseExistingPgSchema=schemaTableNames->Utils.Array.isEmpty,
+      ~cleanRun,
     )
     // Execute all queries within a single transaction for integrity
     let _ = await sql->Postgres.beginSql(sql => {
@@ -461,37 +399,17 @@ let make = (~sql: Postgres.sql, ~pgSchema, ~pgUser): Persistence.storage => {
     })
   }
 
-  let loadEffectCaches = async () => {
-    let schemaTableNames: array<schemaTableName> =
-      await sql->Postgres.unsafe(makeSchemaTableNamesQuery(~pgSchema))
-    schemaTableNames->Belt.Array.keepMapU(schemaTableName => {
-      if schemaTableName.tableName->Js.String2.startsWith("effect_cache_") {
-        Some(
-          (
-            {
-              name: schemaTableName.tableName,
-              size: 0,
-              table: None,
-            }: Persistence.effectCache
-          ),
-        )
-      } else {
-        None
-      }
-    })
-  }
-
   let loadByIdsOrThrow = async (~ids, ~table: Table.table, ~rowsSchema) => {
     switch await (
       switch ids {
       | [_] =>
         sql->Postgres.preparedUnsafe(
-          makeLoadByIdQuery(~pgSchema, ~tableName=table.tableName),
+          makeLoadByIdSql(~pgSchema, ~tableName=table.tableName),
           ids->Obj.magic,
         )
       | _ =>
         sql->Postgres.preparedUnsafe(
-          makeLoadByIdsQuery(~pgSchema, ~tableName=table.tableName),
+          makeLoadByIdsSql(~pgSchema, ~tableName=table.tableName),
           [ids]->Obj.magic,
         )
       }
@@ -516,52 +434,6 @@ let make = (~sql: Postgres.sql, ~pgSchema, ~pgUser): Persistence.storage => {
     }
   }
 
-  let loadByFieldOrThrow = async (
-    ~fieldName: string,
-    ~fieldSchema,
-    ~fieldValue,
-    ~operator: Persistence.operator,
-    ~table: Table.table,
-    ~rowsSchema,
-  ) => {
-    let params = try [fieldValue->S.reverseConvertToJsonOrThrow(fieldSchema)]->Obj.magic catch {
-    | exn =>
-      raise(
-        Persistence.StorageError({
-          message: `Failed loading "${table.tableName}" from storage by field "${fieldName}". Couldn't serialize provided value.`,
-          reason: exn,
-        }),
-      )
-    }
-    switch await sql->Postgres.preparedUnsafe(
-      makeLoadByFieldQuery(
-        ~pgSchema,
-        ~tableName=table.tableName,
-        ~fieldName,
-        ~operator=(operator :> string),
-      ),
-      params,
-    ) {
-    | exception exn =>
-      raise(
-        Persistence.StorageError({
-          message: `Failed loading "${table.tableName}" from storage by field "${fieldName}"`,
-          reason: exn,
-        }),
-      )
-    | rows =>
-      try rows->S.parseOrThrow(rowsSchema) catch {
-      | exn =>
-        raise(
-          Persistence.StorageError({
-            message: `Failed to parse "${table.tableName}" loaded from storage by ids`,
-            reason: exn,
-          }),
-        )
-      }
-    }
-  }
-
   let setOrThrow = (
     type item,
     ~items: array<item>,
@@ -580,8 +452,6 @@ let make = (~sql: Postgres.sql, ~pgSchema, ~pgUser): Persistence.storage => {
   {
     isInitialized,
     initialize,
-    loadByFieldOrThrow,
-    loadEffectCaches,
     loadByIdsOrThrow,
     setOrThrow,
   }