envio 2.26.0-rc.0 → 2.26.0-rc.1

package/src/PgStorage.res

@@ -1,16 +1,16 @@
-let makeCreateIndexSql = (~tableName, ~indexFields, ~pgSchema) => {
+let makeCreateIndexQuery = (~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 makeCreateTableIndicesSql = (table: Table.table, ~pgSchema) => {
+let makeCreateTableIndicesQuery = (table: Table.table, ~pgSchema) => {
   open Belt
   let tableName = table.tableName
   let createIndex = indexField =>
-    makeCreateIndexSql(~tableName, ~indexFields=[indexField], ~pgSchema)
+    makeCreateIndexQuery(~tableName, ~indexFields=[indexField], ~pgSchema)
   let createCompositeIndex = indexFields => {
-    makeCreateIndexSql(~tableName, ~indexFields, ~pgSchema)
+    makeCreateIndexQuery(~tableName, ~indexFields, ~pgSchema)
   }
 
   let singleIndices = table->Table.getSingleIndices
@@ -20,7 +20,7 @@ let makeCreateTableIndicesSql = (table: Table.table, ~pgSchema) => {
     compositeIndices->Array.map(createCompositeIndex)->Js.Array2.joinWith("\n")
 }
 
-let makeCreateTableSql = (table: Table.table, ~pgSchema) => {
+let makeCreateTableQuery = (table: Table.table, ~pgSchema) => {
   open Belt
   let fieldsMapped =
     table
@@ -58,7 +58,7 @@ let makeInitializeTransaction = (
   ~generalTables=[],
   ~entities=[],
   ~enums=[],
-  ~cleanRun=false,
+  ~isEmptyPgSchema=false,
 ) => {
   let allTables = generalTables->Array.copy
   let allEntityTables = []
@@ -71,12 +71,16 @@ let makeInitializeTransaction = (
 
   let query = ref(
     (
-      cleanRun
-        ? `DROP SCHEMA IF EXISTS "${pgSchema}" CASCADE;
-CREATE SCHEMA "${pgSchema}";`
-        : `CREATE SCHEMA IF NOT EXISTS "${pgSchema}";`
+      isEmptyPgSchema && pgSchema === "public"
+      // Hosted Service already have a DB with the created public schema
+      // It also doesn't allow to simply drop it,
+      // so we reuse the existing schema when it's empty
+      // (but only for public, since it's usually always exists)
+        ? ""
+        : `DROP SCHEMA IF EXISTS "${pgSchema}" CASCADE;
+CREATE SCHEMA "${pgSchema}";\n`
     ) ++
-    `GRANT ALL ON SCHEMA "${pgSchema}" TO ${pgUser};
+    `GRANT ALL ON SCHEMA "${pgSchema}" TO "${pgUser}";
 GRANT ALL ON SCHEMA "${pgSchema}" TO public;`,
   )
 
@@ -87,31 +91,17 @@ GRANT ALL ON SCHEMA "${pgSchema}" TO public;`,
       ->Js.Array2.map(v => `'${v->(Utils.magic: Internal.enum => string)}'`)
       ->Js.Array2.joinWith(", ")});`
 
-    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;`
-      }
+    query := query.contents ++ "\n" ++ enumCreateQuery
   })
 
   // Batch all table creation first (optimal for PostgreSQL)
   allTables->Js.Array2.forEach((table: Table.table) => {
-    query := query.contents ++ "\n" ++ makeCreateTableSql(table, ~pgSchema)
+    query := query.contents ++ "\n" ++ makeCreateTableQuery(table, ~pgSchema)
   })
 
   // Then batch all indices (better performance when tables exist)
   allTables->Js.Array2.forEach((table: Table.table) => {
-    let indices = makeCreateTableIndicesSql(table, ~pgSchema)
+    let indices = makeCreateTableIndicesQuery(table, ~pgSchema)
     if indices !== "" {
       query := query.contents ++ "\n" ++ indices
     }
@@ -131,7 +121,7 @@ END IF;`
       query :=
         query.contents ++
         "\n" ++
-        makeCreateIndexSql(
+        makeCreateIndexQuery(
           ~tableName=derivedFromField.derivedFromEntity,
           ~indexFields=[indexField],
           ~pgSchema,
@@ -139,28 +129,40 @@ END IF;`
     })
   })
 
-  [
-    // 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] : [])
+  // Add cache row count function
+  functionsQuery :=
+    functionsQuery.contents ++
+    "\n" ++
+    `CREATE OR REPLACE FUNCTION get_cache_row_count(table_name text) 
+   RETURNS integer AS $$
+   DECLARE
+     result integer;
+   BEGIN
+     EXECUTE format('SELECT COUNT(*) FROM "${pgSchema}".%I', table_name) INTO result;
+     RETURN result;
+   END;
+   $$ LANGUAGE plpgsql;`
+
+  [query.contents]->Js.Array2.concat(
+    functionsQuery.contents !== "" ? [functionsQuery.contents] : [],
+  )
 }
 
-let makeLoadByIdSql = (~pgSchema, ~tableName) => {
+let makeLoadByIdQuery = (~pgSchema, ~tableName) => {
   `SELECT * FROM "${pgSchema}"."${tableName}" WHERE id = $1 LIMIT 1;`
 }
 
-let makeLoadByIdsSql = (~pgSchema, ~tableName) => {
+let makeLoadByFieldQuery = (~pgSchema, ~tableName, ~fieldName, ~operator) => {
+  `SELECT * FROM "${pgSchema}"."${tableName}" WHERE "${fieldName}" ${operator} $1;`
+}
+
+let makeLoadByIdsQuery = (~pgSchema, ~tableName) => {
   `SELECT * FROM "${pgSchema}"."${tableName}" WHERE id = ANY($1::text[]);`
 }
 
-let makeInsertUnnestSetSql = (~pgSchema, ~table: Table.table, ~itemSchema, ~isRawEvents) => {
+let makeInsertUnnestSetQuery = (~pgSchema, ~table: Table.table, ~itemSchema, ~isRawEvents) => {
   let {quotedFieldNames, quotedNonPrimaryFieldNames, arrayFieldTypes} =
-    table->Table.toSqlParams(~schema=itemSchema)
+    table->Table.toSqlParams(~schema=itemSchema, ~pgSchema)
 
   let primaryKeyFieldNames = Table.getPrimaryKeyFieldNames(table)
 
@@ -188,8 +190,9 @@ SELECT * FROM unnest(${arrayFieldTypes
   } ++ ";"
 }
 
-let makeInsertValuesSetSql = (~pgSchema, ~table: Table.table, ~itemSchema, ~itemsCount) => {
-  let {quotedFieldNames, quotedNonPrimaryFieldNames} = table->Table.toSqlParams(~schema=itemSchema)
+let makeInsertValuesSetQuery = (~pgSchema, ~table: Table.table, ~itemSchema, ~itemsCount) => {
+  let {quotedFieldNames, quotedNonPrimaryFieldNames} =
+    table->Table.toSqlParams(~schema=itemSchema, ~pgSchema)
 
   let primaryKeyFieldNames = Table.getPrimaryKeyFieldNames(table)
   let fieldsCount = quotedFieldNames->Array.length
@@ -236,12 +239,13 @@ 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)
+  let {dbSchema, hasArrayField} = table->Table.toSqlParams(~schema=itemSchema, ~pgSchema)
   let isRawEvents = table.tableName === rawEventsTableName
 
   // Should experiment how much it'll affect performance
@@ -253,7 +257,7 @@ let makeTableBatchSetQuery = (~pgSchema, ~table: Table.table, ~itemSchema: S.t<'
 
   if isRawEvents || !hasArrayField {
     {
-      "sql": makeInsertUnnestSetSql(~pgSchema, ~table, ~itemSchema, ~isRawEvents),
+      "query": makeInsertUnnestSetQuery(~pgSchema, ~table, ~itemSchema, ~isRawEvents),
       "convertOrThrow": S.compile(
         S.unnest(dbSchema),
         ~input=Value,
@@ -265,7 +269,12 @@ let makeTableBatchSetQuery = (~pgSchema, ~table: Table.table, ~itemSchema: S.t<'
     }
   } else {
     {
-      "sql": makeInsertValuesSetSql(~pgSchema, ~table, ~itemSchema, ~itemsCount=maxItemsPerQuery),
+      "query": makeInsertValuesSetQuery(
+        ~pgSchema,
+        ~table,
+        ~itemSchema,
+        ~itemsCount=maxItemsPerQuery,
+      ),
       "convertOrThrow": S.compile(
         S.unnest(itemSchema)->S.preprocess(_ => {
           serializer: Utils.Array.flatten->Utils.magic,
@@ -291,6 +300,35 @@ 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) => {
@@ -298,7 +336,7 @@ let setOrThrow = async (sql, ~items, ~table: Table.table, ~itemSchema, ~pgSchema
     ()
   } else {
     // Get or create cached query for this table
-    let query = switch setQueryCache->Utils.WeakMap.get(table) {
+    let data = switch setQueryCache->Utils.WeakMap.get(table) {
     | Some(cached) => cached
     | None => {
         let newQuery = makeTableBatchSetQuery(
@@ -311,50 +349,31 @@ let setOrThrow = async (sql, ~items, ~table: Table.table, ~itemSchema, ~pgSchema
       }
     }
 
-    let sqlQuery = query["sql"]
-
     try {
-      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)
+      if data["isInsertValues"] {
+        let chunks = chunkArray(items, ~chunkSize=maxItemsPerQuery)
         let responses = []
         chunks->Js.Array2.forEach(chunk => {
           let chunkSize = chunk->Array.length
-          let isFullChunk = chunkSize === maxChunkSize
+          let isFullChunk = chunkSize === maxItemsPerQuery
 
           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
-              ? sqlQuery
-              : makeInsertValuesSetSql(
-                  ~pgSchema,
-                  ~table,
-                  ~itemSchema,
-                  ~itemsCount=chunkSize / fieldsCount,
-                ),
-            chunk->Utils.magic,
+              ? data["query"]
+              : makeInsertValuesSetQuery(~pgSchema, ~table, ~itemSchema, ~itemsCount=chunkSize),
+            data["convertOrThrow"](chunk->(Utils.magic: array<'item> => array<unknown>)),
           )
           responses->Js.Array2.push(response)->ignore
         })
         let _ = await Promise.all(responses)
       } else {
         // Use UNNEST approach for single query
-        await sql->Postgres.preparedUnsafe(sqlQuery, payload->Obj.magic)
+        await sql->Postgres.preparedUnsafe(
+          data["query"],
+          data["convertOrThrow"](items->(Utils.magic: array<'item> => array<unknown>)),
+        )
       }
     } catch {
     | S.Raised(_) as exn =>
@@ -368,35 +387,201 @@ let setOrThrow = async (sql, ~items, ~table: Table.table, ~itemSchema, ~pgSchema
       raise(
         Persistence.StorageError({
           message: `Failed to insert items into table "${table.tableName}"`,
-          reason: exn,
+          reason: exn->Internal.prettifyExn,
         }),
       )
     }
   }
 }
 
-let make = (~sql: Postgres.sql, ~pgSchema, ~pgUser): Persistence.storage => {
+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 cacheTablePrefix = "envio_effect_"
+let cacheTablePrefixLength = cacheTablePrefix->String.length
+
+type schemaCacheTableInfo = {
+  @as("table_name")
+  tableName: string,
+  @as("count")
+  count: int,
+}
+
+let makeSchemaCacheTableInfoQuery = (~pgSchema) => {
+  `SELECT 
+    t.table_name,
+    get_cache_row_count(t.table_name) as count
+   FROM information_schema.tables t
+   WHERE t.table_schema = '${pgSchema}' 
+   AND t.table_name LIKE '${cacheTablePrefix}%';`
+}
+
+type psqlExecState =
+  Unknown | Pending(promise<result<string, string>>) | Resolved(result<string, string>)
+
+let getConnectedPsqlExec = {
+  let pgDockerServiceName = "envio-postgres"
+  // Should use the default port, since we're executing the command
+  // from the postgres container's network.
+  let pgDockerServicePort = 5432
+
+  // For development: We run the indexer process locally,
+  //   and there might not be psql installed on the user's machine.
+  //   So we use docker-compose to run psql existing in the postgres container.
+  // For production: We expect indexer to be running in a container,
+  //   with psql installed. So we can call it directly.
+  let psqlExecState = ref(Unknown)
+  async (~pgUser, ~pgHost, ~pgDatabase, ~pgPort) => {
+    switch psqlExecState.contents {
+    | Unknown => {
+        let promise = Promise.make((resolve, _reject) => {
+          let binary = "psql"
+          NodeJs.ChildProcess.exec(`${binary} --version`, (~error, ~stdout as _, ~stderr as _) => {
+            switch error {
+            | Value(_) => {
+                let binary = `docker-compose exec -T -u ${pgUser} ${pgDockerServiceName} psql`
+                NodeJs.ChildProcess.exec(
+                  `${binary} --version`,
+                  (~error, ~stdout as _, ~stderr as _) => {
+                    switch error {
+                    | Value(_) =>
+                      resolve(
+                        Error(`Please check if "psql" binary is installed or docker-compose is running for the local indexer.`),
+                      )
+                    | Null =>
+                      resolve(
+                        Ok(
+                          `${binary} -h ${pgHost} -p ${pgDockerServicePort->Js.Int.toString} -U ${pgUser} -d ${pgDatabase}`,
+                        ),
+                      )
+                    }
+                  },
+                )
+              }
+            | Null =>
+              resolve(
+                Ok(
+                  `${binary} -h ${pgHost} -p ${pgPort->Js.Int.toString} -U ${pgUser} -d ${pgDatabase}`,
+                ),
+              )
+            }
+          })
+        })
+
+        psqlExecState := Pending(promise)
+        let result = await promise
+        psqlExecState := Resolved(result)
+        result
+      }
+    | Pending(promise) => await promise
+    | Resolved(result) => result
+    }
+  }
+}
+
+let make = (
+  ~sql: Postgres.sql,
+  ~pgHost,
+  ~pgSchema,
+  ~pgPort,
+  ~pgUser,
+  ~pgDatabase,
+  ~pgPassword,
+  ~onInitialize=?,
+  ~onNewTables=?,
+): Persistence.storage => {
+  let psqlExecOptions: NodeJs.ChildProcess.execOptions = {
+    env: Js.Dict.fromArray([("PGPASSWORD", pgPassword), ("PATH", %raw(`process.env.PATH`))]),
+  }
+
+  let cacheDirPath = NodeJs.Path.resolve([
+    // Right outside of the generated directory
+    "..",
+    ".envio",
+    "cache",
+  ])
+
   let isInitialized = async () => {
-    let schemas =
+    let envioTables =
       await sql->Postgres.unsafe(
-        `SELECT schema_name FROM information_schema.schemata WHERE schema_name = '${pgSchema}';`,
+        `SELECT table_schema FROM information_schema.tables WHERE table_schema = '${pgSchema}' AND table_name = '${eventSyncStateTableName}';`,
       )
-    schemas->Utils.Array.notEmpty
+    envioTables->Utils.Array.notEmpty
   }
 
-  let initialize = async (~entities=[], ~generalTables=[], ~enums=[], ~cleanRun=false) => {
+  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 queries = makeInitializeTransaction(
       ~pgSchema,
       ~pgUser,
       ~generalTables,
       ~entities,
       ~enums,
-      ~cleanRun,
+      ~isEmptyPgSchema=schemaTableNames->Utils.Array.isEmpty,
     )
     // Execute all queries within a single transaction for integrity
     let _ = await sql->Postgres.beginSql(sql => {
       queries->Js.Array2.map(query => sql->Postgres.unsafe(query))
     })
+
+    switch onInitialize {
+    | Some(onInitialize) => await onInitialize()
+    | None => ()
+    }
   }
 
   let loadByIdsOrThrow = async (~ids, ~table: Table.table, ~rowsSchema) => {
@@ -404,12 +589,12 @@ let make = (~sql: Postgres.sql, ~pgSchema, ~pgUser): Persistence.storage => {
       switch ids {
       | [_] =>
         sql->Postgres.preparedUnsafe(
-          makeLoadByIdSql(~pgSchema, ~tableName=table.tableName),
+          makeLoadByIdQuery(~pgSchema, ~tableName=table.tableName),
           ids->Obj.magic,
         )
       | _ =>
         sql->Postgres.preparedUnsafe(
-          makeLoadByIdsSql(~pgSchema, ~tableName=table.tableName),
+          makeLoadByIdsQuery(~pgSchema, ~tableName=table.tableName),
           [ids]->Obj.magic,
         )
       }
@@ -434,6 +619,52 @@ 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>,
@@ -449,10 +680,212 @@ let make = (~sql: Postgres.sql, ~pgSchema, ~pgUser): Persistence.storage => {
     )
   }
 
+  let setEffectCacheOrThrow = async (
+    ~effectName: string,
+    ~ids: array<string>,
+    ~outputs: array<Internal.effectOutput>,
+    ~outputSchema: S.t<Internal.effectOutput>,
+    ~initialize: bool,
+  ) => {
+    let table = Table.mkTable(
+      cacheTablePrefix ++ effectName,
+      ~fields=[
+        Table.mkField("id", Text, ~fieldSchema=S.string, ~isPrimaryKey=true),
+        Table.mkField("output", JsonB, ~fieldSchema=S.json(~validate=false)),
+      ],
+      ~compositeIndices=[],
+    )
+
+    if initialize {
+      let _ = await sql->Postgres.unsafe(makeCreateTableQuery(table, ~pgSchema))
+      switch onNewTables {
+      | Some(onNewTables) => await onNewTables(~tableNames=[table.tableName])
+      | None => ()
+      }
+    }
+
+    let items = []
+    for idx in 0 to outputs->Array.length - 1 {
+      items
+      ->Js.Array2.push({
+        "id": ids[idx],
+        "output": outputs[idx],
+      })
+      ->ignore
+    }
+
+    await setOrThrow(
+      ~items,
+      ~table,
+      ~itemSchema=S.schema(s =>
+        {
+          "id": s.matches(S.string),
+          "output": s.matches(outputSchema),
+        }
+      ),
+    )
+  }
+
+  let dumpEffectCache = async () => {
+    try {
+      let cacheTableInfo: array<schemaCacheTableInfo> =
+        (await sql
+        ->Postgres.unsafe(makeSchemaCacheTableInfoQuery(~pgSchema)))
+        ->Js.Array2.filter(i => i.count > 0)
+
+      if cacheTableInfo->Utils.Array.notEmpty {
+        // Create .envio/cache directory if it doesn't exist
+        try {
+          await NodeJs.Fs.Promises.access(cacheDirPath)
+        } catch {
+        | _ =>
+          // Create directory if it doesn't exist
+          await NodeJs.Fs.Promises.mkdir(~path=cacheDirPath, ~options={recursive: true})
+        }
+
+        // Command for testing. Run from generated
+        // docker-compose exec -T -u postgres envio-postgres psql -d envio-dev -c 'COPY "public"."envio_effect_getTokenMetadata" TO STDOUT (FORMAT text, HEADER);' > ../.envio/cache/getTokenMetadata.tsv
+
+        switch await getConnectedPsqlExec(~pgUser, ~pgHost, ~pgDatabase, ~pgPort) {
+        | Ok(psqlExec) => {
+            Logging.info(
+              `Dumping cache: ${cacheTableInfo
+                ->Js.Array2.map(({tableName, count}) =>
+                  tableName ++ " (" ++ count->Belt.Int.toString ++ " rows)"
+                )
+                ->Js.Array2.joinWith(", ")}`,
+            )
+
+            let promises = cacheTableInfo->Js.Array2.map(async ({tableName}) => {
+              let cacheName = tableName->Js.String2.sliceToEnd(~from=cacheTablePrefixLength)
+              let outputFile =
+                NodeJs.Path.join(cacheDirPath, cacheName ++ ".tsv")->NodeJs.Path.toString
+
+              let command = `${psqlExec} -c 'COPY "${pgSchema}"."${tableName}" TO STDOUT WITH (FORMAT text, HEADER);' > ${outputFile}`
+
+              Promise.make((resolve, reject) => {
+                NodeJs.ChildProcess.execWithOptions(
+                  command,
+                  psqlExecOptions,
+                  (~error, ~stdout, ~stderr as _) => {
+                    switch error {
+                    | Value(error) => reject(error)
+                    | Null => resolve(stdout)
+                    }
+                  },
+                )
+              })
+            })
+
+            let _ = await promises->Promise.all
+            Logging.info(`Successfully dumped cache to ${cacheDirPath->NodeJs.Path.toString}`)
+          }
+        | Error(message) => Logging.error(`Failed to dump cache. ${message}`)
+        }
+      }
+    } catch {
+    | exn => Logging.errorWithExn(exn->Internal.prettifyExn, `Failed to dump cache.`)
+    }
+  }
+
+  let restoreEffectCache = async (~withUpload) => {
+    if withUpload {
+      // Try to restore cache tables from binary files
+      let nothingToUploadErrorMessage = "Nothing to upload."
+
+      switch await Promise.all2((
+        NodeJs.Fs.Promises.readdir(cacheDirPath)
+        ->Promise.thenResolve(e => Ok(e))
+        ->Promise.catch(_ => Promise.resolve(Error(nothingToUploadErrorMessage))),
+        getConnectedPsqlExec(~pgUser, ~pgHost, ~pgDatabase, ~pgPort),
+      )) {
+      | (Ok(entries), Ok(psqlExec)) => {
+          let cacheFiles = entries->Js.Array2.filter(entry => {
+            entry->Js.String2.endsWith(".tsv")
+          })
+
+          let _ =
+            await cacheFiles
+            ->Js.Array2.map(entry => {
+              let cacheName = entry->Js.String2.slice(~from=0, ~to_=-4) // Remove .tsv extension
+              let tableName = cacheTablePrefix ++ cacheName
+              let table = Table.mkTable(
+                tableName,
+                ~fields=[
+                  Table.mkField("id", Text, ~fieldSchema=S.string, ~isPrimaryKey=true),
+                  Table.mkField("output", JsonB, ~fieldSchema=S.json(~validate=false)),
+                ],
+                ~compositeIndices=[],
+              )
+
+              sql
+              ->Postgres.unsafe(makeCreateTableQuery(table, ~pgSchema))
+              ->Promise.then(() => {
+                let inputFile = NodeJs.Path.join(cacheDirPath, entry)->NodeJs.Path.toString
+
+                let command = `${psqlExec} -c 'COPY "${pgSchema}"."${tableName}" FROM STDIN WITH (FORMAT text, HEADER);' < ${inputFile}`
+
+                Promise.make(
+                  (resolve, reject) => {
+                    NodeJs.ChildProcess.execWithOptions(
+                      command,
+                      psqlExecOptions,
+                      (~error, ~stdout, ~stderr as _) => {
+                        switch error {
+                        | Value(error) => reject(error)
+                        | Null => resolve(stdout)
+                        }
+                      },
+                    )
+                  },
+                )
+              })
+            })
+            ->Promise.all
+
+          Logging.info("Successfully uploaded cache.")
+        }
+      | (Error(message), _)
+      | (_, Error(message)) =>
+        if message === nothingToUploadErrorMessage {
+          Logging.info("No cache found to upload.")
+        } else {
+          Logging.error(`Failed to upload cache, continuing without it. ${message}`)
+        }
+      }
+    }
+
+    let cacheTableInfo: array<schemaCacheTableInfo> =
+      await sql->Postgres.unsafe(makeSchemaCacheTableInfoQuery(~pgSchema))
+
+    if withUpload && cacheTableInfo->Utils.Array.notEmpty {
+      switch onNewTables {
+      | Some(onNewTables) =>
+        await onNewTables(
+          ~tableNames=cacheTableInfo->Js.Array2.map(info => {
+            info.tableName
+          }),
+        )
+      | None => ()
+      }
+    }
+
+    cacheTableInfo->Js.Array2.map((info): Persistence.effectCacheRecord => {
+      {
+        effectName: info.tableName->Js.String2.sliceToEnd(~from=cacheTablePrefixLength),
+        count: info.count,
+      }
+    })
+  }
+
   {
     isInitialized,
     initialize,
+    loadByFieldOrThrow,
     loadByIdsOrThrow,
     setOrThrow,
+    setEffectCacheOrThrow,
+    dumpEffectCache,
+    restoreEffectCache,
   }
 }