feat: add background index creation jobs (#723)

lunary-ai · Jan 23, 2025 · 99985bb · 99985bb
1 parent 85f99ba
commit 99985bb
Show file tree

Hide file tree

Showing 5 changed files with 145 additions and 5 deletions.
diff --git a/packages/backend/src/create-indexes.ts b/packages/backend/src/create-indexes.ts
@@ -0,0 +1,121 @@
+import sql from "./utils/db";
+
+export async function createIndexes() {
+  await sql`select pg_advisory_lock(123456789)`;
+
+  try {
+    const migrations = await sql`
+      select 
+        id, 
+        name, 
+        operation,
+        statement
+      from 
+        _db_migration_index
+       where 
+        status in ('pending', 'failed')
+       order 
+        by id;
+    `;
+
+    for (const migration of migrations) {
+      const { id, name, operation, statement } = migration;
+
+      try {
+        await sql`
+          update 
+            _db_migration_index
+          set 
+            status = 'in-progress'
+           where 
+            id = ${id}
+        `;
+
+        await sql.unsafe(statement);
+
+        if (operation === "create") {
+          const [validCheck] = await sql`
+          select 
+            c.relname, 
+            i.indisvalid
+          from 
+            pg_class c
+            join pg_index i on c.oid = i.indexrelid
+           where 
+            c.relname = ${name}
+            and i.indisvalid = true
+        `;
+
+          if (validCheck) {
+            await sql`
+            update 
+              _db_migration_index
+            set 
+              status = 'done'
+            where 
+              id = ${id}
+          `;
+            console.log(`Index migration "${name}" completed successfully.`);
+          } else {
+            await sql.unsafe(`drop index if exists ${name}`);
+            await sql`
+            update 
+              _db_migration_index
+            set 
+              status = 'failed'
+            where 
+              id = ${id}
+          `;
+            console.warn(
+              `Index migration "${name}" failed; dropped partial index.`,
+            );
+          }
+        } else if (operation === "drop") {
+          const [stillExists] = await sql`
+            select 
+              c.relname
+            from 
+              pg_class c
+              join pg_index i on c.oid = i.indexrelid
+            where c.relname = ${name}
+          `;
+
+          if (!stillExists) {
+            await sql`
+              update 
+                _db_migration_index
+              set 
+                status = 'done'
+              where 
+                id = ${id}
+            `;
+            console.log(`Index drop "${name}" completed successfully.`);
+          } else {
+            await sql`
+              update 
+                _db_migration_index
+              set 
+                status = 'failed'
+              where 
+                id = ${id}
+        `;
+            console.warn(`Index drop "${name}" failed; index still exists.`);
+          }
+        }
+      } catch (err) {
+        console.error(`Index migration "${name}" errored:`, err);
+        await sql.unsafe(`drop index if exists ${name}`);
+        await sql`
+          update 
+            _db_migration_index
+          set 
+            status = 'failed'
+          where 
+            id = ${id}
+        `;
+      }
+    }
+  } finally {
+    await sql`select pg_advisory_unlock(123456789)`;
+  }
+}
diff --git a/packages/backend/src/index.ts b/packages/backend/src/index.ts
@@ -18,11 +18,13 @@ import { initSentry, requestHandler, tracingMiddleWare } from "./utils/sentry";
 import licenseMiddleware from "./utils/license";
 import config from "./utils/config";
 import { startMaterializedViewRefreshJob } from "./jobs/materializedViews";
+import { createIndexes } from "./create-indexes";
 
 checkDbConnection();
 setupCronJobs();
 
 if (process.env.NODE_ENV === "production") {
+  createIndexes();
   startMaterializedViewRefreshJob();
 }
 initSentry();

diff --git a/packages/db/0056.sql b/packages/db/0056.sql
@@ -1,7 +1,10 @@
-drop index if exists run_input_idx;
-drop index if exists run_output_idx;
-drop index if exists run_error_idx;
 drop table if exists log;
 create extension if not exists btree_gin;
-create index run_project_id_input_idx on run using gin (project_id, ((input)::text) gin_trgm_ops);
-create index run_project_id_output_idx on run using gin (project_id, ((output)::text) gin_trgm_ops);
+
+-- !!!WARNING!!! statement below moved to the new index creation system in 0059.sql, but it's kept as a reference for users who have already run this migration
+
+-- drop index if exists run_input_idx;
+-- drop index if exists run_output_idx;
+-- drop index if exists run_error_idx;
+-- create index run_project_id_input_idx on run using gin (project_id, ((input)::text) gin_trgm_ops);
+-- create index run_project_id_output_idx on run using gin (project_id, ((output)::text) gin_trgm_ops);
diff --git a/packages/db/0058.sql b/packages/db/0058.sql
@@ -0,0 +1,8 @@
+create table _db_migration_index (
+  id serial primary key,
+  name text not null,
+  statement text not null,
+  operation text not null default 'create',
+  status text not null default 'pending' -- "pending", "in-progress", "done", "failed"
+);
+
diff --git a/packages/db/0059.sql b/packages/db/0059.sql
@@ -0,0 +1,6 @@
+insert into _db_migration_index (name, operation, statement) values
+('run_input_idx', 'drop', 'drop index concurrently if exists run_input_idx'),
+('run_output_idx', 'drop', 'drop index concurrently if exists run_output_idx'),
+('run_error_idx', 'drop', 'drop index concurrently if exists run_error_idx'),
+('run_project_id_input_idx', 'create', 'create index concurrently if not exists run_project_id_input_idx on run using gin (project_id, ((input)::text) gin_trgm_ops)'),
+('run_project_id_output_idx', 'create', 'create index concurrently if not exists run_project_id_output_idx on run using gin (project_id, ((output)::text) gin_trgm_ops)');