@dbos-inc/dbos-sdk 3.6.5-preview → 3.6.8-preview

package/dist/src/system_database.js

@@ -8,9 +8,6 @@ var __decorate = (this && this.__decorate) || function (decorators, target, key,
 var __metadata = (this && this.__metadata) || function (k, v) {
     if (typeof Reflect === "object" && typeof Reflect.metadata === "function") return Reflect.metadata(k, v);
 };
-var __importDefault = (this && this.__importDefault) || function (mod) {
-    return (mod && mod.__esModule) ? mod : { "default": mod };
-};
 Object.defineProperty(exports, "__esModule", { value: true });
 exports.PostgresSystemDatabase = exports.ensureSystemDatabase = exports.grantDbosSchemaPermissions = exports.DBOS_STREAM_CLOSED_SENTINEL = exports.DBOS_FUNCNAME_CLOSESTREAM = exports.DBOS_FUNCNAME_WRITESTREAM = exports.DBOS_FUNCNAME_GETSTATUS = exports.DBOS_FUNCNAME_SLEEP = exports.DBOS_FUNCNAME_GETEVENT = exports.DBOS_FUNCNAME_SETEVENT = exports.DBOS_FUNCNAME_RECV = exports.DBOS_FUNCNAME_SEND = void 0;
 const dbos_executor_1 = require("./dbos-executor");
@@ -18,7 +15,6 @@ const pg_1 = require("pg");
 const error_1 = require("./error");
 const workflow_1 = require("./workflow");
 const utils_1 = require("./utils");
-const knex_1 = __importDefault(require("knex"));
 const crypto_1 = require("crypto");
 const utils_2 = require("./utils");
 const database_utils_1 = require("./database_utils");
@@ -451,8 +447,6 @@ class PostgresSystemDatabase {
     sysDbPoolSize;
     pool;
     systemPoolConfig;
-    // TODO: remove Knex connection in favor of just using Pool
-    knexDB;
     /*
      * Generally, notifications are asynchronous.  One should:
      *  Subscribe to updates
@@ -498,15 +492,6 @@ class PostgresSystemDatabase {
                 this.logger.warn(`Unexpected error in idle client: ${err}`);
             });
         });
-        const knexConfig = {
-            client: 'pg',
-            connection: this.systemPoolConfig,
-            pool: {
-                min: 0,
-                max: this.sysDbPoolSize,
-            },
-        };
-        this.knexDB = (0, knex_1.default)(knexConfig);
     }
     async init(debugMode = false) {
         await ensureSystemDatabase(this.systemDatabaseUrl, this.logger, debugMode);
@@ -515,7 +500,6 @@ class PostgresSystemDatabase {
         }
     }
     async destroy() {
-        await this.knexDB.destroy();
         if (this.reconnectTimeout) {
             clearTimeout(this.reconnectTimeout);
         }
@@ -1343,41 +1327,65 @@ class PostgresSystemDatabase {
             selectColumns.push('output', 'error');
         }
         input.sortDesc = input.sortDesc ?? false; // By default, sort in ascending order
-        let query = this.knexDB(`${schemaName}.workflow_status`)
-            .select(...selectColumns)
-            .orderBy(`${schemaName}.workflow_status.created_at`, input.sortDesc ? 'desc' : 'asc');
+        // Build WHERE clauses
+        const whereClauses = [];
+        const params = [];
+        let paramCounter = 1;
         if (input.workflowName) {
-            query = query.where(`${schemaName}.workflow_status.name`, input.workflowName);
+            whereClauses.push(`name = $${paramCounter}`);
+            params.push(input.workflowName);
+            paramCounter++;
         }
         if (input.workflow_id_prefix) {
-            query = query.whereLike(`${schemaName}.workflow_status.workflow_uuid`, `${input.workflow_id_prefix}%`);
+            whereClauses.push(`workflow_uuid LIKE $${paramCounter}`);
+            params.push(`${input.workflow_id_prefix}%`);
+            paramCounter++;
         }
         if (input.workflowIDs) {
-            query = query.whereIn(`${schemaName}.workflow_status.workflow_uuid`, input.workflowIDs);
+            const placeholders = input.workflowIDs.map((_, i) => `$${paramCounter + i}`).join(', ');
+            whereClauses.push(`workflow_uuid IN (${placeholders})`);
+            params.push(...input.workflowIDs);
+            paramCounter += input.workflowIDs.length;
         }
         if (input.authenticatedUser) {
-            query = query.where(`${schemaName}.workflow_status.authenticated_user`, input.authenticatedUser);
+            whereClauses.push(`authenticated_user = $${paramCounter}`);
+            params.push(input.authenticatedUser);
+            paramCounter++;
         }
         if (input.startTime) {
-            query = query.where(`${schemaName}.workflow_status.created_at`, '>=', new Date(input.startTime).getTime());
+            whereClauses.push(`created_at >= $${paramCounter}`);
+            params.push(new Date(input.startTime).getTime());
+            paramCounter++;
         }
         if (input.endTime) {
-            query = query.where(`${schemaName}.workflow_status.created_at`, '<=', new Date(input.endTime).getTime());
+            whereClauses.push(`created_at <= $${paramCounter}`);
+            params.push(new Date(input.endTime).getTime());
+            paramCounter++;
         }
         if (input.status) {
-            query = query.where(`${schemaName}.workflow_status.status`, input.status);
+            whereClauses.push(`status = $${paramCounter}`);
+            params.push(input.status);
+            paramCounter++;
         }
         if (input.applicationVersion) {
-            query = query.where(`${schemaName}.workflow_status.application_version`, input.applicationVersion);
-        }
-        if (input.limit) {
-            query = query.limit(input.limit);
-        }
-        if (input.offset) {
-            query = query.offset(input.offset);
-        }
-        const rows = await query;
-        return rows.map(mapWorkflowStatus);
+            whereClauses.push(`application_version = $${paramCounter}`);
+            params.push(input.applicationVersion);
+            paramCounter++;
+        }
+        const whereClause = whereClauses.length > 0 ? `WHERE ${whereClauses.join(' AND ')}` : '';
+        const orderClause = `ORDER BY created_at ${input.sortDesc ? 'DESC' : 'ASC'}`;
+        const limitClause = input.limit ? `LIMIT ${input.limit}` : '';
+        const offsetClause = input.offset ? `OFFSET ${input.offset}` : '';
+        const query = `
+      SELECT ${selectColumns.join(', ')}
+      FROM ${schemaName}.workflow_status
+      ${whereClause}
+      ${orderClause}
+      ${limitClause}
+      ${offsetClause}
+    `;
+        const result = await this.pool.query(query, params);
+        return result.rows.map(mapWorkflowStatus);
     }
     async listQueuedWorkflows(input) {
         const schemaName = dbos_executor_1.DBOSExecutor.systemDBSchemaName;
@@ -1405,34 +1413,54 @@ class PostgresSystemDatabase {
             selectColumns.push('inputs', 'request');
         }
         const sortDesc = input.sortDesc ?? false; // By default, sort in ascending order
-        let query = this.knexDB(`${schemaName}.workflow_status`)
-            .select(...selectColumns)
-            .whereNotNull(`${schemaName}.workflow_status.queue_name`)
-            .whereIn(`${schemaName}.workflow_status.status`, [workflow_1.StatusString.ENQUEUED, workflow_1.StatusString.PENDING])
-            .orderBy(`${schemaName}.workflow_status.created_at`, sortDesc ? 'desc' : 'asc');
+        // Build WHERE clauses
+        const whereClauses = [];
+        const params = [];
+        let paramCounter = 1;
+        // Always filter for queued workflows
+        whereClauses.push(`queue_name IS NOT NULL`);
+        whereClauses.push(`status IN ($${paramCounter}, $${paramCounter + 1})`);
+        params.push(workflow_1.StatusString.ENQUEUED, workflow_1.StatusString.PENDING);
+        paramCounter += 2;
         if (input.workflowName) {
-            query = query.where(`${schemaName}.workflow_status.name`, input.workflowName);
+            whereClauses.push(`name = $${paramCounter}`);
+            params.push(input.workflowName);
+            paramCounter++;
         }
         if (input.queueName) {
-            query = query.where(`${schemaName}.workflow_status.queue_name`, input.queueName);
+            whereClauses.push(`queue_name = $${paramCounter}`);
+            params.push(input.queueName);
+            paramCounter++;
         }
         if (input.startTime) {
-            query = query.where(`${schemaName}.workflow_status.created_at`, '>=', new Date(input.startTime).getTime());
+            whereClauses.push(`created_at >= $${paramCounter}`);
+            params.push(new Date(input.startTime).getTime());
+            paramCounter++;
         }
         if (input.endTime) {
-            query = query.where(`${schemaName}.workflow_status.created_at`, '<=', new Date(input.endTime).getTime());
+            whereClauses.push(`created_at <= $${paramCounter}`);
+            params.push(new Date(input.endTime).getTime());
+            paramCounter++;
         }
         if (input.status) {
-            query = query.where(`${schemaName}.workflow_status.status`, input.status);
-        }
-        if (input.limit) {
-            query = query.limit(input.limit);
-        }
-        if (input.offset) {
-            query = query.offset(input.offset);
-        }
-        const rows = await query;
-        return rows.map(mapWorkflowStatus);
+            whereClauses.push(`status = $${paramCounter}`);
+            params.push(input.status);
+            paramCounter++;
+        }
+        const whereClause = `WHERE ${whereClauses.join(' AND ')}`;
+        const orderClause = `ORDER BY created_at ${sortDesc ? 'DESC' : 'ASC'}`;
+        const limitClause = input.limit ? `LIMIT ${input.limit}` : '';
+        const offsetClause = input.offset ? `OFFSET ${input.offset}` : '';
+        const query = `
+      SELECT ${selectColumns.join(', ')}
+      FROM ${schemaName}.workflow_status
+      ${whereClause}
+      ${orderClause}
+      ${limitClause}
+      ${offsetClause}
+    `;
+        const result = await this.pool.query(query, params);
+        return result.rows.map(mapWorkflowStatus);
     }
     async clearQueueAssignment(workflowID) {
         // Reset the status of the task from "PENDING" to "ENQUEUED"
@@ -1454,19 +1482,20 @@ class PostgresSystemDatabase {
         const startTimeMs = Date.now();
         const limiterPeriodMS = queue.rateLimit ? queue.rateLimit.periodSec * 1000 : 0;
         const claimedIDs = [];
-        await this.knexDB.transaction(async (trx) => {
+        const client = await this.pool.connect();
+        try {
+            await client.query('BEGIN ISOLATION LEVEL REPEATABLE READ');
             // If there is a rate limit, compute how many functions have started in its period.
             let numRecentQueries = 0;
             if (queue.rateLimit) {
-                const numRecentQueriesS = (await trx(`${dbos_executor_1.DBOSExecutor.systemDBSchemaName}.workflow_status`)
-                    .count()
-                    .where('queue_name', queue.name)
-                    .andWhere('status', '<>', workflow_1.StatusString.ENQUEUED)
-                    .andWhere('started_at_epoch_ms', '>', startTimeMs - limiterPeriodMS)
-                    .first()).count;
-                numRecentQueries = parseInt(`${numRecentQueriesS}`);
+                const countResult = await client.query(`SELECT COUNT(*) FROM ${dbos_executor_1.DBOSExecutor.systemDBSchemaName}.workflow_status
+           WHERE queue_name = $1
+             AND status <> $2
+             AND started_at_epoch_ms > $3`, [queue.name, workflow_1.StatusString.ENQUEUED, startTimeMs - limiterPeriodMS]);
+                numRecentQueries = Number(countResult.rows[0].count);
                 if (numRecentQueries >= queue.rateLimit.limitPerPeriod) {
-                    return;
+                    await client.query('COMMIT');
+                    return claimedIDs;
                 }
             }
             // Dequeue functions eligible for this worker and ordered by the time at which they were enqueued.
@@ -1475,15 +1504,12 @@ class PostgresSystemDatabase {
             let maxTasks = Infinity;
             if (queue.workerConcurrency !== undefined || queue.concurrency !== undefined) {
                 // Count how many workflows on this queue are currently PENDING both locally and globally.
-                const runningTasksSubquery = trx(`${dbos_executor_1.DBOSExecutor.systemDBSchemaName}.workflow_status`)
-                    .select('executor_id')
-                    .count('* as task_count')
-                    .where('queue_name', queue.name)
-                    .andWhere('status', workflow_1.StatusString.PENDING)
-                    .groupBy('executor_id');
-                const runningTasksResult = await runningTasksSubquery;
+                const runningTasksResult = await client.query(`SELECT executor_id, COUNT(*) as task_count
+           FROM ${dbos_executor_1.DBOSExecutor.systemDBSchemaName}.workflow_status
+           WHERE queue_name = $1 AND status = $2
+           GROUP BY executor_id`, [queue.name, workflow_1.StatusString.PENDING]);
                 const runningTasksResultDict = {};
-                runningTasksResult.forEach((row) => {
+                runningTasksResult.rows.forEach((row) => {
                     runningTasksResultDict[row.executor_id] = Number(row.task_count);
                 });
                 const runningTasksForThisWorker = runningTasksResultDict[executorID] || 0;
@@ -1501,36 +1527,29 @@ class PostgresSystemDatabase {
             }
             // Return immediately if there are no available tasks due to flow control limits
             if (maxTasks <= 0) {
-                return;
+                await client.query('COMMIT');
+                return claimedIDs;
             }
             // Retrieve the first max_tasks workflows in the queue.
             // Only retrieve workflows of the local version (or without version set)
-            let query = trx(`${dbos_executor_1.DBOSExecutor.systemDBSchemaName}.workflow_status`)
-                .where('status', workflow_1.StatusString.ENQUEUED)
-                .andWhere('queue_name', queue.name)
-                .andWhere((b) => {
-                b.whereNull('application_version').orWhere('application_version', appVersion);
-            })
-                .forUpdate();
-            // Unless global concurrency is set, use skipLocked to only select
-            // rows that can be locked. If global concurrency is set, use noWait
-            // to ensure all processes have a consistent view of the table.
-            if (queue.concurrency) {
-                query = query.noWait();
-            }
-            else {
-                query = query.skipLocked();
-            }
-            if (queue.priorityEnabled) {
-                query = query.orderBy('priority', 'asc').orderBy('created_at', 'asc');
-            }
-            else {
-                query = query.orderBy('created_at', 'asc');
-            }
-            if (maxTasks !== Infinity) {
-                query = query.limit(maxTasks);
-            }
-            const rows = (await query.select(['workflow_uuid']));
+            const lockMode = queue.concurrency ? 'FOR UPDATE NOWAIT' : 'FOR UPDATE SKIP LOCKED';
+            const orderClause = queue.priorityEnabled ? 'ORDER BY priority ASC, created_at ASC' : 'ORDER BY created_at ASC';
+            const limitClause = maxTasks !== Infinity ? `LIMIT ${maxTasks}` : '';
+            const selectQuery = `
+        SELECT workflow_uuid
+        FROM ${dbos_executor_1.DBOSExecutor.systemDBSchemaName}.workflow_status
+        WHERE status = $1
+          AND queue_name = $2
+          AND (application_version IS NULL OR application_version = $3)
+        ${orderClause}
+        ${limitClause}
+        ${lockMode}
+      `;
+            const { rows } = await client.query(selectQuery, [
+                workflow_1.StatusString.ENQUEUED,
+                queue.name,
+                appVersion,
+            ]);
             // Start the workflows
             const workflowIDs = rows.map((row) => row.workflow_uuid);
             for (const id of workflowIDs) {
@@ -1540,18 +1559,28 @@ class PostgresSystemDatabase {
                     break;
                 }
                 // Start the functions by marking them as pending and updating their executor IDs
-                await trx(`${dbos_executor_1.DBOSExecutor.systemDBSchemaName}.workflow_status`)
-                    .where('workflow_uuid', id)
-                    .update({
-                    status: workflow_1.StatusString.PENDING,
-                    executor_id: executorID,
-                    application_version: appVersion,
-                    started_at_epoch_ms: startTimeMs,
-                    workflow_deadline_epoch_ms: trx.raw('CASE WHEN workflow_timeout_ms IS NOT NULL AND workflow_deadline_epoch_ms IS NULL THEN (EXTRACT(epoch FROM now()) * 1000)::bigint + workflow_timeout_ms ELSE workflow_deadline_epoch_ms END'),
-                });
+                await client.query(`UPDATE ${dbos_executor_1.DBOSExecutor.systemDBSchemaName}.workflow_status
+           SET status = $1,
+               executor_id = $2,
+               application_version = $3,
+               started_at_epoch_ms = $4,
+               workflow_deadline_epoch_ms = CASE
+                 WHEN workflow_timeout_ms IS NOT NULL AND workflow_deadline_epoch_ms IS NULL
+                 THEN (EXTRACT(epoch FROM now()) * 1000)::bigint + workflow_timeout_ms
+                 ELSE workflow_deadline_epoch_ms
+               END
+           WHERE workflow_uuid = $5`, [workflow_1.StatusString.PENDING, executorID, appVersion, startTimeMs, id]);
                 claimedIDs.push(id);
             }
-        }, { isolationLevel: 'repeatable read' });
+            await client.query('COMMIT');
+        }
+        catch (error) {
+            await client.query('ROLLBACK');
+            throw error;
+        }
+        finally {
+            client.release();
+        }
         // Return the IDs of all functions we marked started
         return claimedIDs;
     }
@@ -1633,14 +1662,12 @@ class PostgresSystemDatabase {
     async garbageCollect(cutoffEpochTimestampMs, rowsThreshold) {
         if (rowsThreshold !== undefined) {
             // Get the created_at timestamp of the rows_threshold newest row
-            const result = (await this.knexDB(`${dbos_executor_1.DBOSExecutor.systemDBSchemaName}.workflow_status`)
-                .select('created_at')
-                .orderBy('created_at', 'desc')
-                .limit(1)
-                .offset(rowsThreshold - 1)
-                .first());
-            if (result !== undefined) {
-                const rowsBasedCutoff = result.created_at;
+            const result = await this.pool.query(`SELECT created_at
+         FROM ${dbos_executor_1.DBOSExecutor.systemDBSchemaName}.workflow_status
+         ORDER BY created_at DESC
+         LIMIT 1 OFFSET $1`, [rowsThreshold - 1]);
+            if (result.rows.length > 0) {
+                const rowsBasedCutoff = result.rows[0].created_at;
                 // Use the more restrictive cutoff (higher timestamp = more recent = more deletion)
                 if (cutoffEpochTimestampMs === undefined || rowsBasedCutoff > cutoffEpochTimestampMs) {
                     cutoffEpochTimestampMs = rowsBasedCutoff;
@@ -1651,10 +1678,9 @@ class PostgresSystemDatabase {
             return;
         }
         // Delete all workflows older than cutoff that are NOT PENDING or ENQUEUED
-        await this.knexDB(`${dbos_executor_1.DBOSExecutor.systemDBSchemaName}.workflow_status`)
-            .where('created_at', '<', cutoffEpochTimestampMs)
-            .whereNotIn('status', [workflow_1.StatusString.PENDING, workflow_1.StatusString.ENQUEUED])
-            .del();
+        await this.pool.query(`DELETE FROM ${dbos_executor_1.DBOSExecutor.systemDBSchemaName}.workflow_status
+       WHERE created_at < $1
+         AND status NOT IN ($2, $3)`, [cutoffEpochTimestampMs, workflow_1.StatusString.PENDING, workflow_1.StatusString.ENQUEUED]);
         return;
     }
 }