@sixb/pg 0.1.1 → 0.1.2

package/src/pg-execution-storage.ts

@@ -0,0 +1,147 @@
+import {
+  type ExecutionStorageRow,
+  executionRecordFromStorageRow,
+  executionRecordToStorageRow,
+  normalizeExecutionRecord,
+  validateExecutionRecordReferences,
+} from "@sixb/core/internal/execution-storage"
+import type { CreateExecutionInput, ExecutionRecord, ExecutionStorage } from "@sixb/core/storage"
+import { ExecutionStorageError } from "@sixb/core/storage"
+import type { PgAuthStorage } from "./auth-storage"
+import { isUniqueViolation } from "./storage-errors"
+import type { PgStoreClient } from "./transactions"
+
+export class PgExecutionStorage implements ExecutionStorage {
+  constructor(
+    private readonly sql: PgStoreClient,
+    private readonly auth: PgAuthStorage
+  ) {}
+
+  async create(input: CreateExecutionInput): Promise<ExecutionRecord> {
+    const record = normalizeExecutionRecord(input)
+    await validateExecutionRecordReferences(record, {
+      auth: this.auth,
+      getExecution: (params) => this.getById(params),
+    })
+    const row = executionRecordToStorageRow(record)
+
+    try {
+      const [created] = await this.sql<PgExecutionRow[]>`
+        INSERT INTO executions (
+          project_id,
+          id,
+          executor_kind,
+          executor_id,
+          source_kind,
+          source_id,
+          requested_by_user_id,
+          requested_by_service_account_id,
+          correlation_id,
+          parent_execution_id,
+          authority_kind,
+          authority_user_id,
+          authority_service_account_id,
+          authority_session_id,
+          authority_access_token_id,
+          authority_primitive_kind,
+          authority_primitive_id,
+          authority_kernel_operation,
+          created_at
+        ) VALUES (
+          ${row.projectId},
+          ${row.id},
+          ${row.executorKind},
+          ${row.executorId},
+          ${row.sourceKind},
+          ${row.sourceId},
+          ${row.requestedByUserId},
+          ${row.requestedByServiceAccountId},
+          ${row.correlationId},
+          ${row.parentExecutionId},
+          ${row.authorityKind},
+          ${row.authorityUserId},
+          ${row.authorityServiceAccountId},
+          ${row.authoritySessionId},
+          ${row.authorityAccessTokenId},
+          ${row.authorityPrimitiveKind},
+          ${row.authorityPrimitiveId},
+          ${row.authorityKernelOperation},
+          ${row.createdAt}
+        )
+        RETURNING *
+      `
+
+      if (!created) {
+        throw new Error("[SixbPg] Execution insert returned no row.")
+      }
+      return executionRecordFromStorageRow(toStorageRow(created))
+    } catch (error) {
+      if (isUniqueViolation(error)) {
+        throw new ExecutionStorageError(
+          "duplicate_execution",
+          `[SixbPg] Execution '${record.id}' already exists in project '${record.projectId}'.`,
+          { cause: error }
+        )
+      }
+      throw error
+    }
+  }
+
+  async getById(params: {
+    readonly projectId: string
+    readonly id: string
+  }): Promise<ExecutionRecord | null> {
+    const [row] = await this.sql<PgExecutionRow[]>`
+      SELECT *
+      FROM executions
+      WHERE project_id = ${params.projectId} AND id = ${params.id}
+    `
+    return row ? executionRecordFromStorageRow(toStorageRow(row)) : null
+  }
+}
+
+function toStorageRow(row: PgExecutionRow): ExecutionStorageRow {
+  return {
+    projectId: row.project_id,
+    id: row.id,
+    executorKind: row.executor_kind,
+    executorId: row.executor_id,
+    sourceKind: row.source_kind,
+    sourceId: row.source_id,
+    requestedByUserId: row.requested_by_user_id,
+    requestedByServiceAccountId: row.requested_by_service_account_id,
+    correlationId: row.correlation_id,
+    parentExecutionId: row.parent_execution_id,
+    authorityKind: row.authority_kind,
+    authorityUserId: row.authority_user_id,
+    authorityServiceAccountId: row.authority_service_account_id,
+    authoritySessionId: row.authority_session_id,
+    authorityAccessTokenId: row.authority_access_token_id,
+    authorityPrimitiveKind: row.authority_primitive_kind,
+    authorityPrimitiveId: row.authority_primitive_id,
+    authorityKernelOperation: row.authority_kernel_operation,
+    createdAt: new Date(row.created_at),
+  }
+}
+
+interface PgExecutionRow {
+  readonly project_id: string
+  readonly id: string
+  readonly executor_kind: ExecutionStorageRow["executorKind"]
+  readonly executor_id: string
+  readonly source_kind: ExecutionStorageRow["sourceKind"]
+  readonly source_id: string
+  readonly requested_by_user_id: string | null
+  readonly requested_by_service_account_id: string | null
+  readonly correlation_id: string
+  readonly parent_execution_id: string | null
+  readonly authority_kind: ExecutionStorageRow["authorityKind"]
+  readonly authority_user_id: string | null
+  readonly authority_service_account_id: string | null
+  readonly authority_session_id: string | null
+  readonly authority_access_token_id: string | null
+  readonly authority_primitive_kind: ExecutionStorageRow["authorityPrimitiveKind"]
+  readonly authority_primitive_id: string | null
+  readonly authority_kernel_operation: ExecutionStorageRow["authorityKernelOperation"]
+  readonly created_at: Date | string
+}

package/src/pg-pipeline-run-storage.ts

@@ -1,5 +1,9 @@
+import type { JsonValue } from "@sixb/core"
+import { parseSixbFailure, serializeSixbFailure } from "@sixb/core/internal/errors"
+import { assertPipelineRunExecution } from "@sixb/core/internal/pipeline-run-storage-provider"
 import type { DatasetVersionRef } from "@sixb/core/lake-storage"
 import type {
+  ExecutionStorage,
   FinishPipelineRunInput,
   FinishPipelineStepRunInput,
   ListLatestPipelineRunsInput,
@@ -8,14 +12,18 @@ import type {
   ListPipelineRunsResult,
   ListPipelineStepRunsInput,
   ListPipelineStepRunsResult,
-  PipelineRunFailure,
   PipelineRunRecord,
   PipelineRunStorage,
   PipelineStepRunRecord,
+  QueuePipelineRunInput,
   StartPipelineRunInput,
   StartPipelineStepRunInput,
 } from "@sixb/core/storage"
-import { PipelineRunError } from "@sixb/core/storage"
+import {
+  canRequeuePipelineRunAfterEnqueueFailure,
+  PIPELINE_RUN_FAILURE_CODES,
+  PipelineRunError,
+} from "@sixb/core/storage"
 import { queryLatestRunsByOwnerId } from "./latest-run-query"
 import type { SqlParameter } from "./pg-client"
 import { appendRunListFilters, hasEmptyStatuses, queryRunList } from "./run-list-query"
@@ -23,23 +31,37 @@ import { isUniqueViolation } from "./storage-errors"
 import { type PgStoreClient, runPgTransaction } from "./transactions"
 
 export class PgPipelineRunStorage implements PipelineRunStorage {
-  constructor(private readonly sql: PgStoreClient) {}
-
-  async start(input: StartPipelineRunInput): Promise<PipelineRunRecord> {
+  constructor(
+    private readonly sql: PgStoreClient,
+    private readonly executions: ExecutionStorage
+  ) {}
+
+  async queue(input: QueuePipelineRunInput): Promise<PipelineRunRecord> {
+    await assertPipelineRunExecution({
+      executions: this.executions,
+      projectId: input.projectId,
+      executionId: input.executionId,
+      runId: input.id,
+      pipelineId: input.pipelineId,
+    })
     try {
       const [row] = await this.sql<PipelineRunDatabaseRow[]>`
         INSERT INTO pipeline_runs (
           project_id,
           id,
+          execution_id,
           pipeline_id,
           status,
+          queued_at,
           started_at
         ) VALUES (
           ${input.projectId},
           ${input.id},
+          ${input.executionId},
           ${input.pipelineId},
-          ${"running"},
-          ${input.startedAt ?? new Date()}
+          ${"queued"},
+          ${input.queuedAt ?? new Date()},
+          ${null}
         )
         RETURNING *
       `
@@ -47,13 +69,62 @@ export class PgPipelineRunStorage implements PipelineRunStorage {
       return rowToPipelineRunRecord(row)
     } catch (error) {
       if (isUniqueViolation(error)) {
+        return this.requeueAfterEnqueueFailure(input)
+      }
+
+      throw error
+    }
+  }
+
+  private async requeueAfterEnqueueFailure(
+    input: QueuePipelineRunInput
+  ): Promise<PipelineRunRecord> {
+    return runPgTransaction(this.sql, async (tx) => {
+      const [existing] = await tx<PipelineRunDatabaseRow[]>`
+        SELECT * FROM pipeline_runs
+        WHERE project_id = ${input.projectId} AND id = ${input.id}
+        FOR UPDATE
+      `
+      if (
+        !existing ||
+        !canRequeuePipelineRunAfterEnqueueFailure(rowToPipelineRunRecord(existing), input)
+      ) {
         throw new PipelineRunError(
           `[SixbPg] Pipeline run '${input.id}' already exists for project '${input.projectId}'.`
         )
       }
 
-      throw error
+      const [updated] = await tx<PipelineRunDatabaseRow[]>`
+        UPDATE pipeline_runs
+        SET status = ${"queued"}, queued_at = ${input.queuedAt ?? new Date()},
+            started_at = ${null}, finished_at = ${null}, output_dataset_id = ${null},
+            output_version_id = ${null}, error = ${null}
+        WHERE project_id = ${input.projectId} AND id = ${input.id}
+        RETURNING *
+      `
+      return rowToPipelineRunRecord(updated)
+    })
+  }
+
+  async start(input: StartPipelineRunInput): Promise<PipelineRunRecord> {
+    const [updated] = await this.sql<PipelineRunDatabaseRow[]>`
+      UPDATE pipeline_runs
+      SET status = ${"running"}, started_at = ${input.startedAt ?? new Date()},
+          error = ${null}
+      WHERE project_id = ${input.projectId} AND id = ${input.id} AND status = ${"queued"}
+      RETURNING *
+    `
+    if (updated) return rowToPipelineRunRecord(updated)
+
+    const existing = await this.getById({ projectId: input.projectId, id: input.id })
+    if (!existing) {
+      throw new PipelineRunError(
+        `[SixbPg] Pipeline run '${input.id}' not found for project '${input.projectId}'.`
+      )
     }
+    throw new PipelineRunError(
+      `[SixbPg] Pipeline run '${input.id}' cannot start from status '${existing.status}'.`
+    )
   }
 
   async finish(input: FinishPipelineRunInput): Promise<PipelineRunRecord> {
@@ -69,9 +140,14 @@ export class PgPipelineRunStorage implements PipelineRunStorage {
         )
       }
 
-      if (existing.status !== "running") {
+      const enqueueFailure =
+        input.status === "failed" && input.error.code === "queue.enqueue_failed"
+      if (
+        (enqueueFailure && existing.status !== "queued") ||
+        (!enqueueFailure && existing.status !== "running")
+      ) {
         throw new PipelineRunError(
-          `[SixbPg] Pipeline run '${input.id}' for project '${input.projectId}' is already terminal.`
+          `[SixbPg] Pipeline run '${input.id}' cannot finish from status '${existing.status}'.`
         )
       }
 
@@ -84,8 +160,7 @@ export class PgPipelineRunStorage implements PipelineRunStorage {
                 finished_at = ${input.finishedAt ?? new Date()},
                 output_dataset_id = ${input.output?.datasetId ?? null},
                 output_version_id = ${input.output?.versionId ?? null},
-                error_name = ${null},
-                error_message = ${null}
+                error = ${null}
               WHERE project_id = ${input.projectId} AND id = ${input.id}
               RETURNING *
             `
@@ -96,8 +171,7 @@ export class PgPipelineRunStorage implements PipelineRunStorage {
                 finished_at = ${input.finishedAt ?? new Date()},
                 output_dataset_id = ${null},
                 output_version_id = ${null},
-                error_name = ${input.error?.name ?? null},
-                error_message = ${input.error?.message ?? null}
+                error = ${input.error === undefined ? null : serializeSixbFailure(input.error, PIPELINE_RUN_FAILURE_CODES)}::text::jsonb
               WHERE project_id = ${input.projectId} AND id = ${input.id}
               RETURNING *
             `
@@ -208,8 +282,7 @@ export class PgPipelineRunStorage implements PipelineRunStorage {
                 finished_at = ${input.finishedAt ?? new Date()},
                 output_version_id = ${input.output.versionId},
                 rows_written = ${input.rowsWritten ?? existing.rows_written ?? null},
-                error_name = ${null},
-                error_message = ${null}
+                error = ${null}
               WHERE project_id = ${input.projectId} AND id = ${input.id}
               RETURNING *
             `
@@ -220,8 +293,7 @@ export class PgPipelineRunStorage implements PipelineRunStorage {
                 finished_at = ${input.finishedAt ?? new Date()},
                 output_version_id = ${null},
                 rows_written = ${input.rowsWritten ?? existing.rows_written ?? null},
-                error_name = ${input.error?.name ?? null},
-                error_message = ${input.error?.message ?? null}
+                error = ${input.error === undefined ? null : serializeSixbFailure(input.error, PIPELINE_RUN_FAILURE_CODES)}::text::jsonb
               WHERE project_id = ${input.projectId} AND id = ${input.id}
               RETURNING *
             `
@@ -263,7 +335,13 @@ export class PgPipelineRunStorage implements PipelineRunStorage {
       params.push(...input.pipelineIds)
     }
 
-    index = appendRunListFilters(whereClauses, params, index, input)
+    index = appendRunListFilters(
+      whereClauses,
+      params,
+      index,
+      input,
+      "COALESCE(started_at, queued_at)"
+    )
 
     const { rows, total, hasMore } = await queryRunList<PipelineRunDatabaseRow>({
       sql: this.sql,
@@ -356,27 +434,15 @@ export class PgPipelineRunStorage implements PipelineRunStorage {
   }
 }
 
-function toFailure(row: {
-  readonly error_name: string | null
-  readonly error_message: string | null
-}): PipelineRunFailure | undefined {
-  if (!row.error_message) {
-    return undefined
-  }
-
-  return {
-    name: row.error_name ?? undefined,
-    message: row.error_message,
-  }
-}
-
 function rowToPipelineRunRecord(row: PipelineRunDatabaseRow): PipelineRunRecord {
   return {
     id: row.id,
     projectId: row.project_id,
+    executionId: row.execution_id,
     pipelineId: row.pipeline_id,
     status: row.status,
-    startedAt: new Date(row.started_at),
+    queuedAt: new Date(row.queued_at),
+    startedAt: row.started_at ? new Date(row.started_at) : undefined,
     finishedAt: row.finished_at ? new Date(row.finished_at) : undefined,
     output:
       row.output_dataset_id && row.output_version_id
@@ -385,7 +451,7 @@ function rowToPipelineRunRecord(row: PipelineRunDatabaseRow): PipelineRunRecord
             versionId: row.output_version_id,
           }
         : undefined,
-    error: toFailure(row),
+    error: row.error === null ? undefined : parseSixbFailure(row.error, PIPELINE_RUN_FAILURE_CODES),
   }
 }
 
@@ -409,7 +475,7 @@ function rowToPipelineStepRunRecord(row: PipelineStepRunDatabaseRow): PipelineSt
         }
       : undefined,
     rowsWritten: row.rows_written != null ? Number(row.rows_written) : undefined,
-    error: toFailure(row),
+    error: row.error === null ? undefined : parseSixbFailure(row.error, PIPELINE_RUN_FAILURE_CODES),
   }
 }
 
@@ -428,14 +494,15 @@ function assertOptionalNonNegativeInteger(value: number | undefined, fieldName:
 interface PipelineRunDatabaseRow {
   project_id: string
   id: string
+  execution_id: string
   pipeline_id: string
   status: PipelineRunRecord["status"]
-  started_at: Date | string
+  queued_at: Date | string
+  started_at: Date | string | null
   finished_at: Date | string | null
   output_dataset_id: string | null
   output_version_id: string | null
-  error_name: string | null
-  error_message: string | null
+  error: JsonValue | null
 }
 
 interface PipelineStepRunDatabaseRow {
@@ -452,6 +519,5 @@ interface PipelineStepRunDatabaseRow {
   inputs: readonly DatasetVersionRef[] | string
   output_version_id: string | null
   rows_written: number | string | null
-  error_name: string | null
-  error_message: string | null
+  error: JsonValue | null
 }