@dbos-inc/dbos-enterprise 5.4.1-preview → 5.5.0

package/dist/src/sharding/sys_db.js

@@ -0,0 +1,776 @@
+// Copyright (c) 2026 DBOS, Inc. All rights reserved.
+//
+// This file is part of Transact Enterprise, proprietary software licensed under the
+// DBOS Conductor Software License: https://www.dbos.dev/conductor-license
+// You may not use, copy, modify, or distribute this file except as that license permits.
+// Disabling or circumventing the entitlement checks in @dbos-inc/dbos-sdk's src/enterprise.ts
+// that gate it is a material breach of that license.
+"use strict";
+var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
+    if (k2 === undefined) k2 = k;
+    var desc = Object.getOwnPropertyDescriptor(m, k);
+    if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
+      desc = { enumerable: true, get: function() { return m[k]; } };
+    }
+    Object.defineProperty(o, k2, desc);
+}) : (function(o, m, k, k2) {
+    if (k2 === undefined) k2 = k;
+    o[k2] = m[k];
+}));
+var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
+    Object.defineProperty(o, "default", { enumerable: true, value: v });
+}) : function(o, v) {
+    o["default"] = v;
+});
+var __importStar = (this && this.__importStar) || function (mod) {
+    if (mod && mod.__esModule) return mod;
+    var result = {};
+    if (mod != null) for (var k in mod) if (k !== "default" && Object.prototype.hasOwnProperty.call(mod, k)) __createBinding(result, mod, k);
+    __setModuleDefault(result, mod);
+    return result;
+};
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.ShardedSystemDatabase = void 0;
+const context_1 = require("@dbos-inc/dbos-sdk/internal/context");
+const error_1 = require("@dbos-inc/dbos-sdk/internal/error");
+const debugpoint_1 = require("@dbos-inc/dbos-sdk/internal/debugpoint");
+const serialization_1 = require("@dbos-inc/dbos-sdk/internal/serialization");
+const shard_keys_1 = require("@dbos-inc/dbos-sdk/internal/shard_keys");
+const system_database_1 = require("@dbos-inc/dbos-sdk/internal/system_database");
+const utils_1 = require("@dbos-inc/dbos-sdk/internal/utils");
+const workflow_1 = require("@dbos-inc/dbos-sdk/internal/workflow");
+const migrations = __importStar(require("./migrations"));
+const schema_1 = require("./schema");
+const CITUS_COMMIT_CONFLICT = 'could not commit transaction for shard';
+function quoteIdentifier(name) {
+    return `"${name.replace(/"/g, '""')}"`;
+}
+class ShardedSystemDatabase extends system_database_1.SystemDatabase {
+    discoveryIntervalMs = 60_000;
+    discoveryEmptyIntervalMs = 5_000;
+    claimSlotsPerKey = 8;
+    dbPollingIntervalEventMs = 1_000;
+    #onCitus = undefined;
+    #discovered = new Map();
+    #rediscovering = new Set();
+    #foundWork = new Map();
+    constructor(options) {
+        super(options.systemDatabaseUrl, options.logger, options.serializer, options.sysDbPoolSize ?? system_database_1.DEFAULT_POOL_SIZE, options.systemDatabasePool, options.schemaName, false, options.pollingConcurrency, undefined, options.appName, options.observabilityQueryTimeoutMs, options.idleTransactionTimeoutMs);
+    }
+    get sharded() {
+        return true;
+    }
+    resolvedKey(workflowID, explicit) {
+        if (explicit !== undefined && explicit !== null) {
+            (0, shard_keys_1.validateShardKey)(explicit);
+            return explicit;
+        }
+        const bound = (0, shard_keys_1.boundShardKey)(workflowID);
+        if (bound !== undefined) {
+            return bound;
+        }
+        const ctx = (0, context_1.getCurrentContextStore)();
+        if (ctx !== undefined && ctx.workflowId === workflowID && ctx.shardKey) {
+            return ctx.shardKey;
+        }
+        return (0, shard_keys_1.scopedShardKey)();
+    }
+    shardKeyOf(workflowID, explicit) {
+        const key = this.resolvedKey(workflowID, explicit);
+        if (key !== undefined) {
+            return key;
+        }
+        throw new error_1.DBOSError(`No shard key for workflow ${workflowID}. On a sharded system database, pass the workflow's shardKey, ` +
+            'or use a handle that carries it.');
+    }
+    scopeKey() {
+        const scoped = (0, shard_keys_1.scopedShardKey)();
+        if (scoped === undefined) {
+            throw new error_1.DBOSError('No shard key for this operation. On a sharded system database, set one with DBOS.withShardKey(...).');
+        }
+        return scoped;
+    }
+    singleKey(keys, operation) {
+        const distinct = new Set(keys);
+        if (distinct.size > 1) {
+            throw new error_1.DBOSError(`On a sharded system database, one ${operation} acts only on workflows under a single shard key; ` +
+                `${operation} each key's workflows separately`);
+        }
+        return [...distinct][0];
+    }
+    groupByKey(workflowIDs) {
+        const groups = new Map();
+        for (const workflowID of workflowIDs) {
+            const key = this.shardKeyOf(workflowID);
+            const ids = groups.get(key) ?? [];
+            ids.push(workflowID);
+            groups.set(key, ids);
+        }
+        return groups;
+    }
+    wf(alias, workflowID, params, column = 'workflow_uuid', shardKey) {
+        const key = this.shardKeyOf(workflowID, shardKey);
+        params.push(key, workflowID);
+        return `${(0, system_database_1.qualify)(alias, 'shard_key')} = $${params.length - 1} AND ${(0, system_database_1.qualify)(alias, column)} = $${params.length}`;
+    }
+    wfIn(alias, workflowIDs, params) {
+        const groups = this.groupByKey(workflowIDs);
+        if (groups.size === 0) {
+            params.push(workflowIDs);
+            return `${(0, system_database_1.qualify)(alias, 'workflow_uuid')} = ANY($${params.length}::text[])`;
+        }
+        const predicates = Array.from(groups, ([key, ids]) => {
+            params.push(key, ids);
+            return `${(0, system_database_1.qualify)(alias, 'shard_key')} = $${params.length - 1} AND ${(0, system_database_1.qualify)(alias, 'workflow_uuid')} = ANY($${params.length}::text[])`;
+        });
+        if (predicates.length === 1) {
+            return predicates[0];
+        }
+        return `(${predicates.map((p) => `(${p})`).join(' OR ')})`;
+    }
+    wfJoin(left, right) {
+        return `${left}.shard_key = ${right}.shard_key AND ${left}.workflow_uuid = ${right}.workflow_uuid`;
+    }
+    wfValues(workflowID, shardKey) {
+        return { shard_key: this.shardKeyOf(workflowID, shardKey) };
+    }
+    parentValues(parentWorkflowID) {
+        if (parentWorkflowID === undefined) {
+            return {};
+        }
+        return { parent_shard_key: this.shardKeyOf(parentWorkflowID) };
+    }
+    childValues(childWorkflowID) {
+        if (childWorkflowID === undefined || childWorkflowID === null) {
+            return {};
+        }
+        return { child_shard_key: this.resolvedKey(childWorkflowID) ?? null };
+    }
+    wfColumns(workflowID, params) {
+        params.push(this.shardKeyOf(workflowID));
+        return [['shard_key', `$${params.length}::text`]];
+    }
+    keyConflict(...columns) {
+        return ['shard_key', ...columns].join(', ');
+    }
+    async isCitus() {
+        if (this.#onCitus === undefined) {
+            const client = await this.connect();
+            try {
+                this.#onCitus = await migrations.citusInstalled(client);
+            }
+            finally {
+                client.release();
+            }
+        }
+        return this.#onCitus;
+    }
+    async backgroundJobLock(jobName, callback) {
+        const name = `dbos-background-job:${this.schemaName}:${this.appName ?? ''}:${jobName}`;
+        const claim = `/* pgdog_sharding_key: "${schema_1.GLOBAL_SHARD_KEY}" */ ` +
+            `SELECT ${quoteIdentifier(this.schemaName)}.claim_gate(${(0, utils_1.quoteLiteral)(name)}, 1)`;
+        return await this.lockedBackgroundJob(jobName, `${claim} FROM set_config('idle_in_transaction_session_timeout', '${Math.trunc(this.backgroundJobLockIdleTimeoutMs)}', true)`, async (client) => {
+            const { rows } = await client.query({ text: claim, rowMode: 'array' });
+            if (rows[0]?.[0] !== true) {
+                throw new error_1.DBOSError(`Lost the ${jobName} lock`);
+            }
+        }, callback);
+    }
+    async outcomeInOneStatement() {
+        return !(await this.isCitus());
+    }
+    async insertRows(client, table, columns, rows, suffix = '') {
+        const out = [];
+        for (const row of rows) {
+            out.push(...(await super.insertRows(client, table, columns, [row], suffix)));
+        }
+        return out;
+    }
+    claimGate(queueName, partitionKey, params) {
+        const name = `dbos-claim:${queueName}:${this.scopeKey()}:${partitionKey ?? ''}`;
+        params.push(name, this.claimSlotsPerKey);
+        return `(SELECT ${quoteIdentifier(this.schemaName)}.claim_gate($${params.length - 1}, $${params.length}::int4))`;
+    }
+    claimLockQuery(queueName, partitionKey, params) {
+        const name = JSON.stringify(['dbos-claim-lock', this.schemaName, queueName, this.scopeKey(), partitionKey ?? null]);
+        params.push(name, queueName, workflow_1.StatusString.ENQUEUED);
+        const n = params.length;
+        const s = quoteIdentifier(this.schemaName);
+        return `SELECT (SELECT ${s}.claim_gate($${n - 2}, 1)) AS acquired
+       FROM ${s}.workflow_status
+       WHERE ${this.scoped(undefined, params)} AND queue_name = $${n - 1} AND status = $${n}
+       LIMIT 1`;
+    }
+    partitionLockQuery(queueName, partitionKeys, params) {
+        const name = JSON.stringify(['dbos-partition-claim', this.schemaName, queueName, this.scopeKey()]);
+        params.push(partitionKeys, name, queueName, workflow_1.StatusString.ENQUEUED);
+        const n = params.length;
+        const s = quoteIdentifier(this.schemaName);
+        return `SELECT ARRAY(
+              SELECT pk FROM unnest($${n - 3}::text[]) AS pk
+              WHERE ${s}.claim_gate($${n - 2}::text || ':' || pk, 1)
+            ) AS held
+       FROM ${s}.workflow_status
+       WHERE ${this.scoped(undefined, params)} AND queue_name = $${n - 1} AND status = $${n}
+       LIMIT 1`;
+    }
+    scoped(alias, params) {
+        params.push(this.scopeKey());
+        return `${(0, system_database_1.qualify)(alias, 'shard_key')} = $${params.length}`;
+    }
+    scopedColumns(params) {
+        params.push(this.scopeKey());
+        return [['shard_key', `$${params.length}::text`]];
+    }
+    global(alias, params) {
+        params.push(schema_1.GLOBAL_SHARD_KEY);
+        return `${(0, system_database_1.qualify)(alias, 'shard_key')} = $${params.length}`;
+    }
+    globalValues() {
+        return { shard_key: schema_1.GLOBAL_SHARD_KEY };
+    }
+    keyIn(alias, shardKeys, params, column = 'shard_key') {
+        for (const shardKey of shardKeys) {
+            (0, shard_keys_1.validateShardKey)(shardKey);
+        }
+        if (shardKeys.length === 0) {
+            return 'FALSE';
+        }
+        const binds = shardKeys.map((shardKey) => `$${params.push(shardKey)}`);
+        return `${(0, system_database_1.qualify)(alias, column)} IN (${binds.join(', ')})`;
+    }
+    keyColumn(alias, column = 'shard_key') {
+        return (0, system_database_1.qualify)(alias, column);
+    }
+    keyHint(workflowID) {
+        return this.shardKeyOf(workflowID);
+    }
+    scheduleValues(schedule) {
+        const shardKey = schedule.shardKey;
+        if (shardKey === undefined || shardKey === null) {
+            throw new error_1.DBOSError(`Schedule '${schedule.scheduleName}' has no shard key. On a sharded system database, every schedule ` +
+                'names the shard key its workflows run under.');
+        }
+        (0, shard_keys_1.validateShardKey)(shardKey);
+        return { workflow_shard_key: shardKey };
+    }
+    scheduleKeyColumn() {
+        return 'workflow_shard_key';
+    }
+    async runMigrations() {
+        await migrations.runMigrations(this);
+    }
+    async verifyMigrations() {
+        await migrations.verifyMigrations(this);
+    }
+    isSerializationConflict(e) {
+        return super.isSerializationConflict(e) || (e instanceof Error && e.message.includes(CITUS_COMMIT_CONFLICT));
+    }
+    static requireKey(status) {
+        if (status.shardKey === undefined || status.shardKey === null) {
+            throw new error_1.DBOSError(`Workflow ${status.workflowUUID} has no shard key. On a sharded system database, start or enqueue every ` +
+                'workflow under an explicit shard key: with DBOS.withShardKey(...), or with the shardKey enqueue option ' +
+                'for DBOSClient.');
+        }
+        (0, shard_keys_1.validateShardKey)(status.shardKey);
+    }
+    async insertWorkflowStatus(client, initStatus, creatorXid) {
+        ShardedSystemDatabase.requireKey(initStatus);
+        if (initStatus.deduplicationID !== undefined && initStatus.deduplicationID !== null) {
+            throw new error_1.DBOSError('Deduplication is not supported on a sharded system database: a deduplication ID could only be unique ' +
+                'within one shard key.');
+        }
+        return await super.insertWorkflowStatus(client, initStatus, creatorXid);
+    }
+    debounceDelayedWorkflow() {
+        throw new error_1.DBOSError('Debouncing is not supported on a sharded system database: it relies on deduplication, which could only be ' +
+            'unique within one shard key.');
+    }
+    enqueueWorkflows() {
+        throw new error_1.DBOSError('Batch-inserting workflows is not supported on a sharded system database.');
+    }
+    async enqueueScheduledWorkflows(statuses, firedAt) {
+        for (const status of statuses) {
+            ShardedSystemDatabase.requireKey(status);
+        }
+        const byKey = new Map();
+        for (const status of statuses) {
+            const keyed = byKey.get(status.shardKey) ?? [];
+            keyed.push(status);
+            byKey.set(status.shardKey, keyed);
+        }
+        const inserted = new Set();
+        try {
+            for (const keyed of byKey.values()) {
+                const ids = await (0, system_database_1.withDbRetry)(async () => {
+                    const createdAts = await this.assignBatchCreatedAt(keyed);
+                    return await this.#transaction((client) => this.insertWorkflowBatch(client, keyed, createdAts));
+                });
+                for (const id of ids) {
+                    inserted.add(id);
+                }
+            }
+        }
+        finally {
+            if (inserted.size > 0) {
+                await (0, system_database_1.withDbRetry)(() => this.#transaction((client) => this.stampLastFiredAt(client, statuses, firedAt, inserted)));
+            }
+        }
+        return inserted;
+    }
+    renameApplication() {
+        throw new error_1.DBOSError('Renaming an application is not supported on a sharded system database.');
+    }
+    async precheckOwner(workflowID) {
+        if (workflowID === undefined) {
+            return;
+        }
+        const ownerXid = (0, context_1.currentOwnerXid)(workflowID);
+        if (ownerXid !== undefined) {
+            await (0, system_database_1.withDbRetry)(async () => {
+                const client = await this.connect();
+                try {
+                    await this.inTransaction(client, () => this.checkOwner(client, workflowID, ownerXid));
+                }
+                finally {
+                    client.release();
+                }
+            });
+        }
+    }
+    async initChildWorkflowStatus(initStatus, creatorXid, parentWorkflowID, parentFunctionID, startTimeEpochMs, endTimeEpochMs, reusePolicy = 'return-existing') {
+        if (reusePolicy === 'reject') {
+            throw new error_1.DBOSError("On a sharded system database, workflowIDReusePolicy 'reject' is not supported for a workflow started " +
+                'from inside another workflow; use it only outside workflows.');
+        }
+        ShardedSystemDatabase.requireKey(initStatus);
+        const sameKey = initStatus.shardKey === this.shardKeyOf(parentWorkflowID);
+        if (!sameKey && initStatus.status === workflow_1.StatusString.PENDING) {
+            throw new error_1.DBOSError(`Workflow ${parentWorkflowID} can't start or call child workflow ${initStatus.workflowUUID} under another ` +
+                "shard key. On a sharded system database, a child runs under its parent's key unless it's enqueued.");
+        }
+        const parentOwnerXid = (0, context_1.currentOwnerXid)(parentWorkflowID);
+        const result = await (0, system_database_1.withDbRetry)(async () => {
+            if (sameKey) {
+                return await this.#transaction(async (client) => {
+                    if (parentOwnerXid !== undefined) {
+                        await this.checkOwner(client, parentWorkflowID, parentOwnerXid);
+                    }
+                    const inserted = await this.initWorkflowStatusInternal(client, initStatus, creatorXid, reusePolicy);
+                    await this.recordChild(client, initStatus, parentWorkflowID, parentFunctionID, startTimeEpochMs, endTimeEpochMs);
+                    return inserted;
+                });
+            }
+            if (parentOwnerXid !== undefined) {
+                await this.#transaction((client) => this.checkOwner(client, parentWorkflowID, parentOwnerXid));
+            }
+            const inserted = await this.#transaction((client) => this.initWorkflowStatusInternal(client, initStatus, creatorXid, reusePolicy));
+            await this.#transaction(async (client) => {
+                if (parentOwnerXid !== undefined) {
+                    await this.checkOwner(client, parentWorkflowID, parentOwnerXid);
+                }
+                await this.recordChild(client, initStatus, parentWorkflowID, parentFunctionID, startTimeEpochMs, endTimeEpochMs);
+            });
+            return inserted;
+        });
+        await (0, debugpoint_1.debugTriggerPoint)(debugpoint_1.DEBUG_TRIGGER_INITWF_COMMIT);
+        return result;
+    }
+    async #transaction(fn) {
+        const client = await this.connect();
+        try {
+            return await this.inTransaction(client, () => fn(client));
+        }
+        finally {
+            client.release();
+        }
+    }
+    async recordChild(client, status, parentWorkflowID, parentFunctionID, startTimeEpochMs, endTimeEpochMs) {
+        const childWorkflowID = status.workflowUUID;
+        const values = {
+            ...this.wfValues(parentWorkflowID),
+            workflow_uuid: parentWorkflowID,
+            function_id: parentFunctionID,
+            function_name: status.workflowName,
+            child_workflow_id: childWorkflowID,
+            child_shard_key: status.shardKey,
+            started_at_epoch_ms: startTimeEpochMs,
+            completed_at_epoch_ms: endTimeEpochMs,
+            application_name: this.appName ?? null,
+        };
+        const columns = Object.keys(values);
+        const { rows: inserted } = await client.query(`INSERT INTO "${this.schemaName}".operation_outputs (${columns.join(', ')}, retention_timestamp)
+       VALUES (${columns.map((_, i) => `$${i + 1}`).join(', ')}, (EXTRACT(EPOCH FROM now()) * 1000)::bigint)
+       ON CONFLICT DO NOTHING
+       RETURNING function_id`, Object.values(values));
+        if (inserted.length === 0) {
+            const params = [parentFunctionID];
+            const { rows: existing } = await client.query(`SELECT child_workflow_id FROM "${this.schemaName}".operation_outputs
+         WHERE ${this.wf(undefined, parentWorkflowID, params)} AND function_id = $1`, params);
+            if (existing.length === 0 || existing[0].child_workflow_id !== childWorkflowID) {
+                throw new error_1.DBOSStepNondeterminismError(parentWorkflowID, parentFunctionID);
+            }
+        }
+    }
+    keyedMessage(destinationID, shardKey, idempotencyKey, workflowID, functionID) {
+        if (shardKey === undefined || shardKey === null) {
+            throw new error_1.DBOSError(`Message to workflow ${destinationID} has no shard key. On a sharded system database, pass the ` +
+                "destination's shardKey.");
+        }
+        (0, shard_keys_1.validateShardKey)(shardKey);
+        if (workflowID !== undefined && !idempotencyKey) {
+            const senderKey = this.shardKeyOf(workflowID);
+            idempotencyKey = `${[...senderKey].length}:${senderKey}:${workflowID}::${functionID}::0`;
+        }
+        return { key: shardKey, idempotencyKey };
+    }
+    async #deliver(key, db, destinationID, message, topic, serialization, idempotencyKey) {
+        await (0, shard_keys_1.shardKeyScope)(key, () => this.sendDirectInternal(db, destinationID, message, topic, serialization, idempotencyKey, key));
+    }
+    async send(workflowID, functionID, destinationID, message, topic, serialization, idempotencyKey, shardKey) {
+        const startTime = Date.now();
+        const recorded = await this.getOperationResultAndThrowIfCancelled(workflowID, functionID);
+        if (recorded !== undefined) {
+            if (recorded.functionName !== system_database_1.DBOS_FUNCNAME_SEND) {
+                throw new error_1.DBOSUnexpectedStepError(workflowID, functionID, system_database_1.DBOS_FUNCNAME_SEND, recorded.functionName);
+            }
+            if (recorded.error) {
+                throw await (0, serialization_1.deserializeResError)(recorded.error, recorded.serialization ?? null, this.serializer);
+            }
+            return;
+        }
+        const keyed = this.keyedMessage(destinationID, shardKey, idempotencyKey, workflowID, functionID);
+        await this.precheckOwner(workflowID);
+        try {
+            await (0, system_database_1.withDbRetry)(() => this.#deliver(keyed.key, this.pool, destinationID, message, topic, serialization, keyed.idempotencyKey));
+        }
+        catch (e) {
+            if (e instanceof error_1.DBOSNonExistentWorkflowError) {
+                await this.recordOperationError(workflowID, functionID, system_database_1.DBOS_FUNCNAME_SEND, startTime, e);
+            }
+            throw e;
+        }
+        await this.recordOperationResult(workflowID, functionID, system_database_1.DBOS_FUNCNAME_SEND, true, startTime, Date.now(), {
+            output: null,
+            serialization: serialization_1.DBOSPortableJSON.name(),
+        });
+    }
+    async sendDirect(destinationID, message, topic, serialization, idempotencyKey, client, shardKey) {
+        const keyed = this.keyedMessage(destinationID, shardKey, idempotencyKey, undefined, undefined);
+        if (client !== undefined) {
+            await this.#deliver(keyed.key, client, destinationID, message, topic, serialization, keyed.idempotencyKey);
+            return;
+        }
+        await (0, system_database_1.withDbRetry)(() => this.#deliver(keyed.key, this.pool, destinationID, message, topic, serialization, keyed.idempotencyKey));
+    }
+    async sendFromStep(workflowID, destinationID, message, topic, serialization, idempotencyKey, shardKey) {
+        const keyed = this.keyedMessage(destinationID, shardKey, idempotencyKey, undefined, undefined);
+        await this.precheckOwner(workflowID);
+        await (0, system_database_1.withDbRetry)(() => this.#deliver(keyed.key, this.pool, destinationID, message, topic, serialization, keyed.idempotencyKey));
+        await this.precheckOwner(workflowID);
+    }
+    async runTransactionalStep(workflowID, functionID, functionName, callback) {
+        const startTime = Date.now();
+        const recorded = await this.getOperationResultAndThrowIfCancelled(workflowID, functionID);
+        if (recorded !== undefined) {
+            return recorded;
+        }
+        let output;
+        try {
+            output = await (0, system_database_1.withDbRetry)(() => this.#transaction(callback));
+        }
+        catch (e) {
+            await this.recordOperationError(workflowID, functionID, functionName, startTime, e);
+            throw e;
+        }
+        await this.recordOperationResult(workflowID, functionID, functionName, true, startTime, Date.now(), { output });
+        return undefined;
+    }
+    async getWorkflowStatus(workflowID, callerID, callerFN) {
+        if (!callerID || callerFN === undefined) {
+            return await super.getWorkflowStatus(workflowID);
+        }
+        const startTime = Date.now();
+        const recorded = await this.getOperationResultAndThrowIfCancelled(callerID, callerFN);
+        if (recorded !== undefined) {
+            if (recorded.functionName !== system_database_1.DBOS_FUNCNAME_GETSTATUS) {
+                throw new error_1.DBOSUnexpectedStepError(callerID, callerFN, system_database_1.DBOS_FUNCNAME_GETSTATUS, recorded.functionName);
+            }
+            if (recorded.error) {
+                throw await (0, serialization_1.deserializeResError)(recorded.error, recorded.serialization ?? null, this.serializer);
+            }
+            return recorded.output ? JSON.parse(recorded.output) : null;
+        }
+        const status = await super.getWorkflowStatus(workflowID);
+        await this.recordOperationResult(callerID, callerFN, system_database_1.DBOS_FUNCNAME_GETSTATUS, true, startTime, Date.now(), {
+            output: status ? JSON.stringify(status) : null,
+            serialization: serialization_1.DBOSPortableJSON.name(),
+        });
+        return status;
+    }
+    keyOfBatch(workflowIDs, operation) {
+        return this.singleKey(workflowIDs.map((id) => this.shardKeyOf(id)), operation);
+    }
+    async deadLetterWorkflows(workflowIDs, minRecoveryAttempts) {
+        const key = this.keyOfBatch(workflowIDs, 'deadLetterWorkflows');
+        if (key === undefined) {
+            return;
+        }
+        await (0, shard_keys_1.shardKeyScope)(key, () => super.deadLetterWorkflows(workflowIDs, minRecoveryAttempts));
+    }
+    async resumeWorkflows(workflowIDs, queueName) {
+        const key = this.keyOfBatch(workflowIDs, 'resumeWorkflows');
+        if (key === undefined) {
+            return;
+        }
+        await (0, shard_keys_1.shardKeyScope)(key, () => super.resumeWorkflows(workflowIDs, queueName));
+    }
+    async listWorkflows(input, options = {}) {
+        const keys = input.shardKey;
+        if (typeof keys === 'string' || (Array.isArray(keys) && keys.length === 1)) {
+            return await super.listWorkflows(input, options);
+        }
+        const limit = input.limit;
+        const offset = input.offset ?? 0;
+        const rows = await super.listWorkflows({ ...input, limit: limit ? offset + limit : undefined, offset: undefined }, options);
+        const desc = input.sortDesc ?? false;
+        rows.sort((a, b) => {
+            const byCreated = (a.createdAt ?? 0) - (b.createdAt ?? 0);
+            const order = byCreated !== 0 ? byCreated : a.workflowUUID < b.workflowUUID ? -1 : a.workflowUUID > b.workflowUUID ? 1 : 0;
+            return desc ? -order : order;
+        });
+        return rows.slice(offset, limit ? offset + limit : undefined);
+    }
+    async getWorkflowStatuses(workflowIDs) {
+        const found = new Map();
+        for (const [key, ids] of this.groupByKey(workflowIDs)) {
+            const statuses = await (0, shard_keys_1.shardKeyScope)(key, () => super.getWorkflowStatuses(ids));
+            for (const [id, status] of statuses) {
+                found.set(id, status);
+            }
+        }
+        return new Map(workflowIDs.filter((id) => found.has(id)).map((id) => [id, found.get(id)]));
+    }
+    async checkFirstWorkflowId(workflowIds) {
+        if (workflowIds.length === 0) {
+            throw new error_1.DBOSError('workflowIds must not be empty');
+        }
+        for (const [key, ids] of this.groupByKey(workflowIds)) {
+            const first = await (0, shard_keys_1.shardKeyScope)(key, () => super.checkFirstWorkflowId(ids));
+            if (first !== undefined) {
+                return first;
+            }
+        }
+        return undefined;
+    }
+    async checkCompletedWorkflowIds(workflowIds) {
+        const completed = [];
+        for (const [key, ids] of this.groupByKey(workflowIds)) {
+            completed.push(...(await (0, shard_keys_1.shardKeyScope)(key, () => super.checkCompletedWorkflowIds(ids))));
+        }
+        return completed;
+    }
+    async childrenOf(parents) {
+        if (parents.length === 0) {
+            return [];
+        }
+        const byKey = ShardedSystemDatabase.#pairsByKey(parents);
+        const params = [];
+        const predicates = Array.from(byKey, ([key, ids]) => {
+            params.push(key, ids);
+            return `(parent_shard_key = $${params.length - 1} AND parent_workflow_id = ANY($${params.length}::text[]))`;
+        });
+        const { rows } = await (0, system_database_1.withDbRetry)(() => this.pool.query(`SELECT workflow_uuid, shard_key FROM "${this.schemaName}".workflow_status WHERE ${predicates.join(' OR ')}`, params));
+        return rows.map((row) => [row.workflow_uuid, row.shard_key]);
+    }
+    getDirectChildren() {
+        throw new error_1.DBOSError('On a sharded system database, children are found by (workflow ID, shard key) with childrenOf');
+    }
+    static #pairsByKey(pairs) {
+        const groups = new Map();
+        for (const [workflowID, key] of pairs) {
+            const ids = groups.get(key) ?? [];
+            ids.push(workflowID);
+            groups.set(key, ids);
+        }
+        return groups;
+    }
+    async #applyByLevel(workflowIDs, descend, apply, operation) {
+        const rootKey = this.keyOfBatch(workflowIDs, operation);
+        const pairKey = (id, key) => `${key.length}:${key}:${id}`;
+        let level = workflowIDs.map((id) => [id, rootKey]);
+        const visited = new Set(level.map(([id, key]) => pairKey(id, key)));
+        while (level.length > 0) {
+            for (const [key, ids] of ShardedSystemDatabase.#pairsByKey(level)) {
+                await (0, shard_keys_1.bindShardKeys)(ids.map((id) => [id, key]), () => (0, shard_keys_1.shardKeyScope)(key, () => apply(ids)));
+            }
+            if (!descend) {
+                return;
+            }
+            level = (await this.childrenOf(level)).filter(([id, key]) => !visited.has(pairKey(id, key)));
+            for (const [id, key] of level) {
+                visited.add(pairKey(id, key));
+            }
+        }
+    }
+    async cancelWorkflows(workflowIDs, cancelChildren = false) {
+        await this.#applyByLevel(workflowIDs, cancelChildren, (ids) => super.cancelWorkflows(ids, false), 'cancelWorkflows');
+    }
+    async deleteWorkflows(workflowIDs, deleteChildren = false) {
+        await this.#applyByLevel(workflowIDs, deleteChildren, (ids) => this.deleteWorkflowsInternal(ids), 'deleteWorkflows');
+    }
+    async bulkForkWorkflows(originalWorkflowIDs, forkedWorkflowIDs, startSteps, options = {}) {
+        if (originalWorkflowIDs.length !== forkedWorkflowIDs.length || originalWorkflowIDs.length !== startSteps.length) {
+            throw new error_1.DBOSError('originalWorkflowIDs, forkedWorkflowIDs, and startSteps must have the same length');
+        }
+        const key = this.keyOfBatch(originalWorkflowIDs, 'forkWorkflow');
+        if (key === undefined) {
+            return forkedWorkflowIDs;
+        }
+        await (0, shard_keys_1.shardKeyScope)(key, () => (0, shard_keys_1.bindShardKeys)(forkedWorkflowIDs.map((f) => [f, key]), () => super.bulkForkWorkflows(originalWorkflowIDs, forkedWorkflowIDs, startSteps, options)));
+        return forkedWorkflowIDs;
+    }
+    async dequeueScopes(queueName) {
+        const cached = this.#discovered.get(queueName);
+        const found = this.#foundWork.get(queueName);
+        this.#foundWork.delete(queueName);
+        const now = performance.now();
+        if (cached !== undefined &&
+            cached.keys.length === 0 &&
+            now - cached.at < Math.min(this.discoveryEmptyIntervalMs, this.discoveryIntervalMs) * cached.share) {
+            return [];
+        }
+        if (cached !== undefined && cached.keys.length > 0 && this.discoveryIntervalMs > 0) {
+            if (found === false || now - cached.at >= this.discoveryIntervalMs * cached.share) {
+                this.#rediscover(queueName, cached.keys);
+            }
+            return [...cached.keys];
+        }
+        const scopes = await this.discoverScopes(queueName);
+        this.#remember(queueName, scopes);
+        return [...scopes];
+    }
+    #noteClaim(queueName, found) {
+        this.#foundWork.set(queueName, (this.#foundWork.get(queueName) ?? false) || found);
+    }
+    async claimStartableWorkflows(queue, executorID, appVersion, queuePartitionKey, localRunningCount = 0, partitionLocalRunningCount = 0, ownerXid = null, isLatestVersion) {
+        const claim = await super.claimStartableWorkflows(queue, executorID, appVersion, queuePartitionKey, localRunningCount, partitionLocalRunningCount, ownerXid, isLatestVersion);
+        this.#noteClaim(queue.name, claim.claimedIDs.length > 0 || claim.stopReason !== undefined);
+        return claim;
+    }
+    async findAndMarkStartablePartitionedWorkflows(queue, executorID, appVersion, maxTasks = Infinity, ownerXid = null, isLatestVersion) {
+        const claimed = await super.findAndMarkStartablePartitionedWorkflows(queue, executorID, appVersion, maxTasks, ownerXid, isLatestVersion);
+        this.#noteClaim(queue.name, claimed.length > 0);
+        return claimed;
+    }
+    async getQueuePartitions(queueName) {
+        const partitions = await super.getQueuePartitions(queueName);
+        if (partitions.length === 0) {
+            this.#noteClaim(queueName, false);
+        }
+        return partitions;
+    }
+    #remember(queueName, keys) {
+        this.#discovered.set(queueName, { at: performance.now(), share: 0.8 + Math.random() * 0.4, keys });
+    }
+    #rediscover(queueName, keys) {
+        if (this.#rediscovering.has(queueName)) {
+            return;
+        }
+        this.#rediscovering.add(queueName);
+        void (async () => {
+            try {
+                this.#remember(queueName, await this.discoverScopes(queueName));
+            }
+            catch (e) {
+                this.logger.warn(`Exception finding work for queue ${queueName}: ${e.message}`);
+                this.#remember(queueName, keys);
+            }
+            finally {
+                this.#rediscovering.delete(queueName);
+            }
+        })();
+    }
+    async discoverScopes(queueName) {
+        const params = [queueName];
+        const enqueued = (alias) => `${alias}.queue_name = $1 AND ${alias}.status = '${workflow_1.StatusString.ENQUEUED}' AND ${this.appNameFilter(`${alias}.application_name`, this.appName, params)}`;
+        if (await this.isCitus()) {
+            const filter = enqueued('ws');
+            const { rows } = await this.pool.query(`SELECT DISTINCT ws.shard_key FROM "${this.schemaName}".workflow_status ws WHERE ${filter}`, params);
+            return rows.map((row) => row.shard_key);
+        }
+        return await this.#allKeys(enqueued('ws'), params);
+    }
+    async #allKeys(enqueued, params) {
+        const { rows } = await this.pool.query(`WITH RECURSIVE keys(k) AS (
+         SELECT ''::text
+         UNION ALL
+         SELECT (SELECT min(ws.shard_key) FROM "${this.schemaName}".workflow_status ws
+                 WHERE ${enqueued} AND ws.shard_key > keys.k)
+         FROM keys WHERE keys.k IS NOT NULL
+       )
+       SELECT k FROM keys WHERE k IS NOT NULL AND k != ''`, params);
+        return rows.map((row) => row.k);
+    }
+    async cancelTimedOutWorkflows(limit) {
+        const nowMs = Date.now();
+        const active = [workflow_1.StatusString.ENQUEUED, workflow_1.StatusString.PENDING, workflow_1.StatusString.DELAYED, workflow_1.StatusString.SUSPENDED];
+        const params = [active, nowMs, limit];
+        const scope = this.appNameFilter('application_name', this.appName, params);
+        const { rows: candidates } = await this.pool.query(`SELECT workflow_uuid, shard_key FROM "${this.schemaName}".workflow_status
+       WHERE status = ANY($1::text[])
+         AND workflow_deadline_epoch_ms IS NOT NULL
+         AND workflow_deadline_epoch_ms <= $2
+         AND ${scope}
+       ORDER BY workflow_deadline_epoch_ms
+       LIMIT $3`, params);
+        const cancelled = [];
+        for (const [key, ids] of ShardedSystemDatabase.#pairsByKey(candidates.map((r) => [r.workflow_uuid, r.shard_key]))) {
+            const { rows } = await this.pool.query(`UPDATE "${this.schemaName}".workflow_status
+         SET status = $1, owner_xid = NULL, queue_name = NULL, started_at_epoch_ms = NULL, suspend_key = NULL,
+             updated_at = (EXTRACT(EPOCH FROM now()) * 1000)::bigint,
+             completed_at = (EXTRACT(EPOCH FROM now()) * 1000)::bigint
+         WHERE shard_key = $2 AND workflow_uuid = ANY($3::text[])
+           AND status = ANY($4::text[]) AND workflow_deadline_epoch_ms <= $5
+         RETURNING workflow_uuid`, [workflow_1.StatusString.CANCELLED, key, ids, active, nowMs]);
+            cancelled.push(...rows.map((row) => row.workflow_uuid));
+        }
+        return cancelled;
+    }
+    async reenqueueWorkflowsForRecovery(executorID, appVersion, recoveryQueueName) {
+        const params = [workflow_1.StatusString.PENDING, executorID, appVersion];
+        const scope = this.appNameFilter('application_name', this.appName, params);
+        const { rows: pending } = await (0, system_database_1.withDbRetry)(() => this.pool.query(`SELECT workflow_uuid, shard_key FROM "${this.schemaName}".workflow_status
+         WHERE status = $1 AND executor_id = $2 AND application_version = $3 AND ${scope}`, params));
+        const recovered = [];
+        for (const [key, ids] of ShardedSystemDatabase.#pairsByKey(pending.map((r) => [r.workflow_uuid, r.shard_key]))) {
+            const updateParams = [
+                workflow_1.StatusString.ENQUEUED,
+                recoveryQueueName,
+                workflow_1.StatusString.PENDING,
+                executorID,
+                key,
+                ids,
+            ];
+            let rows;
+            try {
+                ({ rows } = await (0, system_database_1.withDbRetry)(() => this.pool.query(`UPDATE "${this.schemaName}".workflow_status
+           SET started_at_epoch_ms = NULL,
+               status = $1,
+               owner_xid = NULL,
+               updated_at = (EXTRACT(EPOCH FROM now()) * 1000)::bigint,
+               queue_name = COALESCE(queue_name, $2)
+           WHERE shard_key = $5 AND workflow_uuid = ANY($6::text[]) AND status = $3 AND executor_id = $4
+           RETURNING workflow_uuid`, updateParams)));
+            }
+            catch (e) {
+                this.logger.error(`Exception encountered when recovering the workflows of shard key ${key}: ${String(e)}`);
+                continue;
+            }
+            recovered.push(...rows.map((row) => ({ workflowID: row.workflow_uuid, shardKey: key })));
+        }
+        return recovered;
+    }
+}
+exports.ShardedSystemDatabase = ShardedSystemDatabase;
+//# sourceMappingURL=sys_db.js.map