@mastra/pg 0.3.1-alpha.2 → 0.3.1-alpha.4

package/src/storage/index.test.ts

@@ -4,7 +4,7 @@ import type { MessageType } from '@mastra/core/memory';
 import { TABLE_WORKFLOW_SNAPSHOT, TABLE_MESSAGES, TABLE_THREADS, TABLE_EVALS } from '@mastra/core/storage';
 import type { WorkflowRunState } from '@mastra/core/workflows';
 import pgPromise from 'pg-promise';
-import { describe, it, expect, beforeAll, beforeEach, afterAll, afterEach } from 'vitest';
+import { describe, it, expect, beforeAll, beforeEach, afterAll, afterEach, vi } from 'vitest';
 
 import { PostgresStore } from '.';
 import type { PostgresConfig } from '.';
@@ -19,6 +19,8 @@ const TEST_CONFIG: PostgresConfig = {
 
 const connectionString = `postgresql://${TEST_CONFIG.user}:${TEST_CONFIG.password}@${TEST_CONFIG.host}:${TEST_CONFIG.port}/${TEST_CONFIG.database}`;
 
+vi.setConfig({ testTimeout: 60_000, hookTimeout: 60_000 });
+
 // Sample test data factory functions
 const createSampleThread = () => ({
   id: `thread-${randomUUID()}`,
@@ -29,17 +31,20 @@ const createSampleThread = () => ({
   metadata: { key: 'value' },
 });
 
-const createSampleMessage = (threadId: string) =>
-  ({
-    id: `msg-${randomUUID()}`,
-    role: 'user',
-    type: 'text',
-    threadId,
-    content: [{ type: 'text', text: 'Hello' }],
-    createdAt: new Date(),
-  }) as any;
-
-const createSampleWorkflowSnapshot = (status: string, createdAt?: Date) => {
+const createSampleMessage = (threadId: string): MessageType => ({
+  id: `msg-${randomUUID()}`,
+  resourceId: `resource-${randomUUID()}`,
+  role: 'user',
+  type: 'text',
+  threadId,
+  content: [{ type: 'text', text: 'Hello' }],
+  createdAt: new Date(),
+});
+
+const createSampleWorkflowSnapshot = (
+  status: WorkflowRunState['context']['steps'][string]['status'],
+  createdAt?: Date,
+) => {
   const runId = `run-${randomUUID()}`;
   const stepId = `step-${randomUUID()}`;
   const timestamp = createdAt || new Date();
@@ -58,9 +63,10 @@ const createSampleWorkflowSnapshot = (status: string, createdAt?: Date) => {
       attempts: {},
     },
     activePaths: [],
+    suspendedPaths: {},
     runId,
     timestamp: timestamp.getTime(),
-  } as WorkflowRunState;
+  };
   return { snapshot, runId, stepId };
 };
 
@@ -82,6 +88,13 @@ const createSampleEval = (agentName: string, isTest = false) => {
   };
 };
 
+const checkWorkflowSnapshot = (snapshot: WorkflowRunState | string, stepId: string, status: string) => {
+  if (typeof snapshot === 'string') {
+    throw new Error('Expected WorkflowRunState, got string');
+  }
+  expect(snapshot.context?.steps[stepId]?.status).toBe(status);
+};
+
 describe('PostgresStore', () => {
   let store: PostgresStore;
 
@@ -233,7 +246,11 @@ describe('PostgresStore', () => {
       // Retrieve messages
       const retrievedMessages = await store.getMessages({ threadId: thread.id });
       expect(retrievedMessages).toHaveLength(2);
-      expect(retrievedMessages).toEqual(expect.arrayContaining(messages));
+      const checkMessages = messages.map(m => {
+        const { resourceId, ...rest } = m;
+        return rest;
+      });
+      expect(retrievedMessages).toEqual(expect.arrayContaining(checkMessages));
     });
 
     it('should handle empty message array', async () => {
@@ -253,12 +270,13 @@ describe('PostgresStore', () => {
 
       await store.saveMessages({ messages });
 
-      const retrievedMessages = await store.getMessages({ threadId: thread.id });
+      const retrievedMessages = await store.getMessages<MessageType>({ threadId: thread.id });
       expect(retrievedMessages).toHaveLength(3);
 
       // Verify order is maintained
       retrievedMessages.forEach((msg, idx) => {
-        expect((msg.content[0] as any).text).toBe((messages[idx].content[0] as any).text);
+        // @ts-expect-error
+        expect(msg.content[0].text).toBe(messages[idx].content[0].text);
       });
     });
 
@@ -338,11 +356,17 @@ describe('PostgresStore', () => {
       const snapshot = {
         status: 'running',
         context: {
+          steps: {},
           stepResults: {},
           attempts: {},
           triggerData: { type: 'manual' },
         },
-      } as any;
+        value: {},
+        activePaths: [],
+        suspendedPaths: {},
+        runId,
+        timestamp: new Date().getTime(),
+      };
 
       await store.persistWorkflowSnapshot({
         workflowName,
@@ -373,28 +397,40 @@ describe('PostgresStore', () => {
       const initialSnapshot = {
         status: 'running',
         context: {
+          steps: {},
           stepResults: {},
           attempts: {},
           triggerData: { type: 'manual' },
         },
+        value: {},
+        activePaths: [],
+        suspendedPaths: {},
+        runId,
+        timestamp: new Date().getTime(),
       };
 
       await store.persistWorkflowSnapshot({
         workflowName,
         runId,
-        snapshot: initialSnapshot as any,
+        snapshot: initialSnapshot,
       });
 
       const updatedSnapshot = {
         status: 'completed',
         context: {
+          steps: {},
           stepResults: {
             'step-1': { status: 'success', result: { data: 'test' } },
           },
           attempts: { 'step-1': 1 },
           triggerData: { type: 'manual' },
         },
-      } as any;
+        value: {},
+        activePaths: [],
+        suspendedPaths: {},
+        runId,
+        timestamp: new Date().getTime(),
+      };
 
       await store.persistWorkflowSnapshot({
         workflowName,
@@ -432,6 +468,7 @@ describe('PostgresStore', () => {
               dependencies: ['step-3', 'step-4'],
             },
           },
+          steps: {},
           attempts: { 'step-1': 1, 'step-2': 0 },
           triggerData: {
             type: 'scheduled',
@@ -453,6 +490,7 @@ describe('PostgresStore', () => {
             status: 'waiting',
           },
         ],
+        suspendedPaths: {},
         runId: runId,
         timestamp: Date.now(),
       };
@@ -460,7 +498,7 @@ describe('PostgresStore', () => {
       await store.persistWorkflowSnapshot({
         workflowName,
         runId,
-        snapshot: complexSnapshot as unknown as WorkflowRunState,
+        snapshot: complexSnapshot,
       });
 
       const loadedSnapshot = await store.loadWorkflowSnapshot({
@@ -486,8 +524,8 @@ describe('PostgresStore', () => {
       const workflowName1 = 'default_test_1';
       const workflowName2 = 'default_test_2';
 
-      const { snapshot: workflow1, runId: runId1, stepId: stepId1 } = createSampleWorkflowSnapshot('completed');
-      const { snapshot: workflow2, runId: runId2, stepId: stepId2 } = createSampleWorkflowSnapshot('running');
+      const { snapshot: workflow1, runId: runId1, stepId: stepId1 } = createSampleWorkflowSnapshot('success');
+      const { snapshot: workflow2, runId: runId2, stepId: stepId2 } = createSampleWorkflowSnapshot('failed');
 
       await store.persistWorkflowSnapshot({ workflowName: workflowName1, runId: runId1, snapshot: workflow1 });
       await new Promise(resolve => setTimeout(resolve, 10)); // Small delay to ensure different timestamps
@@ -498,17 +536,17 @@ describe('PostgresStore', () => {
       expect(total).toBe(2);
       expect(runs[0]!.workflowName).toBe(workflowName2); // Most recent first
       expect(runs[1]!.workflowName).toBe(workflowName1);
-      const firstSnapshot = runs[0]!.snapshot as WorkflowRunState;
-      const secondSnapshot = runs[1]!.snapshot as WorkflowRunState;
-      expect(firstSnapshot.context?.steps[stepId2]?.status).toBe('running');
-      expect(secondSnapshot.context?.steps[stepId1]?.status).toBe('completed');
+      const firstSnapshot = runs[0]!.snapshot;
+      const secondSnapshot = runs[1]!.snapshot;
+      checkWorkflowSnapshot(firstSnapshot, stepId2, 'failed');
+      checkWorkflowSnapshot(secondSnapshot, stepId1, 'success');
     });
 
     it('filters by workflow name', async () => {
       const workflowName1 = 'filter_test_1';
       const workflowName2 = 'filter_test_2';
 
-      const { snapshot: workflow1, runId: runId1, stepId: stepId1 } = createSampleWorkflowSnapshot('completed');
+      const { snapshot: workflow1, runId: runId1, stepId: stepId1 } = createSampleWorkflowSnapshot('success');
       const { snapshot: workflow2, runId: runId2 } = createSampleWorkflowSnapshot('failed');
 
       await store.persistWorkflowSnapshot({ workflowName: workflowName1, runId: runId1, snapshot: workflow1 });
@@ -519,8 +557,8 @@ describe('PostgresStore', () => {
       expect(runs).toHaveLength(1);
       expect(total).toBe(1);
       expect(runs[0]!.workflowName).toBe(workflowName1);
-      const snapshot = runs[0]!.snapshot as WorkflowRunState;
-      expect(snapshot.context?.steps[stepId1]?.status).toBe('completed');
+      const snapshot = runs[0]!.snapshot;
+      checkWorkflowSnapshot(snapshot, stepId1, 'success');
     });
 
     it('filters by date range', async () => {
@@ -531,9 +569,9 @@ describe('PostgresStore', () => {
       const workflowName2 = 'date_test_2';
       const workflowName3 = 'date_test_3';
 
-      const { snapshot: workflow1, runId: runId1 } = createSampleWorkflowSnapshot('completed');
-      const { snapshot: workflow2, runId: runId2, stepId: stepId2 } = createSampleWorkflowSnapshot('running');
-      const { snapshot: workflow3, runId: runId3, stepId: stepId3 } = createSampleWorkflowSnapshot('waiting');
+      const { snapshot: workflow1, runId: runId1 } = createSampleWorkflowSnapshot('success');
+      const { snapshot: workflow2, runId: runId2, stepId: stepId2 } = createSampleWorkflowSnapshot('failed');
+      const { snapshot: workflow3, runId: runId3, stepId: stepId3 } = createSampleWorkflowSnapshot('suspended');
 
       await store.insert({
         tableName: TABLE_WORKFLOW_SNAPSHOT,
@@ -574,10 +612,10 @@ describe('PostgresStore', () => {
       expect(runs).toHaveLength(2);
       expect(runs[0]!.workflowName).toBe(workflowName3);
       expect(runs[1]!.workflowName).toBe(workflowName2);
-      const firstSnapshot = runs[0]!.snapshot as WorkflowRunState;
-      const secondSnapshot = runs[1]!.snapshot as WorkflowRunState;
-      expect(firstSnapshot.context?.steps[stepId3]?.status).toBe('waiting');
-      expect(secondSnapshot.context?.steps[stepId2]?.status).toBe('running');
+      const firstSnapshot = runs[0]!.snapshot;
+      const secondSnapshot = runs[1]!.snapshot;
+      checkWorkflowSnapshot(firstSnapshot, stepId3, 'suspended');
+      checkWorkflowSnapshot(secondSnapshot, stepId2, 'failed');
     });
 
     it('handles pagination', async () => {
@@ -585,9 +623,9 @@ describe('PostgresStore', () => {
       const workflowName2 = 'page_test_2';
       const workflowName3 = 'page_test_3';
 
-      const { snapshot: workflow1, runId: runId1, stepId: stepId1 } = createSampleWorkflowSnapshot('completed');
-      const { snapshot: workflow2, runId: runId2, stepId: stepId2 } = createSampleWorkflowSnapshot('running');
-      const { snapshot: workflow3, runId: runId3, stepId: stepId3 } = createSampleWorkflowSnapshot('waiting');
+      const { snapshot: workflow1, runId: runId1, stepId: stepId1 } = createSampleWorkflowSnapshot('success');
+      const { snapshot: workflow2, runId: runId2, stepId: stepId2 } = createSampleWorkflowSnapshot('failed');
+      const { snapshot: workflow3, runId: runId3, stepId: stepId3 } = createSampleWorkflowSnapshot('suspended');
 
       await store.persistWorkflowSnapshot({ workflowName: workflowName1, runId: runId1, snapshot: workflow1 });
       await new Promise(resolve => setTimeout(resolve, 10)); // Small delay to ensure different timestamps
@@ -601,18 +639,119 @@ describe('PostgresStore', () => {
       expect(page1.total).toBe(3); // Total count of all records
       expect(page1.runs[0]!.workflowName).toBe(workflowName3);
       expect(page1.runs[1]!.workflowName).toBe(workflowName2);
-      const firstSnapshot = page1.runs[0]!.snapshot as WorkflowRunState;
-      const secondSnapshot = page1.runs[1]!.snapshot as WorkflowRunState;
-      expect(firstSnapshot.context?.steps[stepId3]?.status).toBe('waiting');
-      expect(secondSnapshot.context?.steps[stepId2]?.status).toBe('running');
+      const firstSnapshot = page1.runs[0]!.snapshot;
+      const secondSnapshot = page1.runs[1]!.snapshot;
+      checkWorkflowSnapshot(firstSnapshot, stepId3, 'suspended');
+      checkWorkflowSnapshot(secondSnapshot, stepId2, 'failed');
 
       // Get second page
       const page2 = await store.getWorkflowRuns({ limit: 2, offset: 2 });
       expect(page2.runs).toHaveLength(1);
       expect(page2.total).toBe(3);
       expect(page2.runs[0]!.workflowName).toBe(workflowName1);
-      const snapshot = page2.runs[0]!.snapshot as WorkflowRunState;
-      expect(snapshot.context?.steps[stepId1]?.status).toBe('completed');
+      const snapshot = page2.runs[0]!.snapshot;
+      checkWorkflowSnapshot(snapshot, stepId1, 'success');
+    });
+  });
+
+  describe('getWorkflowRunById', () => {
+    const workflowName = 'workflow-id-test';
+    let runId: string;
+    let stepId: string;
+
+    beforeEach(async () => {
+      // Insert a workflow run for positive test
+      const sample = createSampleWorkflowSnapshot('success');
+      runId = sample.runId;
+      stepId = sample.stepId;
+      await store.insert({
+        tableName: TABLE_WORKFLOW_SNAPSHOT,
+        record: {
+          workflow_name: workflowName,
+          run_id: runId,
+          resourceId: 'resource-abc',
+          snapshot: sample.snapshot,
+          createdAt: new Date(),
+          updatedAt: new Date(),
+        },
+      });
+    });
+
+    it('should retrieve a workflow run by ID', async () => {
+      const found = await store.getWorkflowRunById({
+        runId,
+        workflowName,
+      });
+      expect(found).not.toBeNull();
+      expect(found?.runId).toBe(runId);
+      checkWorkflowSnapshot(found?.snapshot!, stepId, 'success');
+    });
+
+    it('should return null for non-existent workflow run ID', async () => {
+      const notFound = await store.getWorkflowRunById({
+        runId: 'non-existent-id',
+        workflowName,
+      });
+      expect(notFound).toBeNull();
+    });
+  });
+  describe('getWorkflowRuns with resourceId', () => {
+    const workflowName = 'workflow-id-test';
+    let resourceId: string;
+    let runIds: string[] = [];
+
+    beforeEach(async () => {
+      // Insert multiple workflow runs for the same resourceId
+      resourceId = 'resource-shared';
+      for (const status of ['success', 'failed']) {
+        const sample = createSampleWorkflowSnapshot(status as WorkflowRunState['context']['steps'][string]['status']);
+        runIds.push(sample.runId);
+        await store.insert({
+          tableName: TABLE_WORKFLOW_SNAPSHOT,
+          record: {
+            workflow_name: workflowName,
+            run_id: sample.runId,
+            resourceId,
+            snapshot: sample.snapshot,
+            createdAt: new Date(),
+            updatedAt: new Date(),
+          },
+        });
+      }
+      // Insert a run with a different resourceId
+      const other = createSampleWorkflowSnapshot('waiting');
+      await store.insert({
+        tableName: TABLE_WORKFLOW_SNAPSHOT,
+        record: {
+          workflow_name: workflowName,
+          run_id: other.runId,
+          resourceId: 'resource-other',
+          snapshot: other.snapshot,
+          createdAt: new Date(),
+          updatedAt: new Date(),
+        },
+      });
+    });
+
+    it('should retrieve all workflow runs by resourceId', async () => {
+      const { runs } = await store.getWorkflowRuns({
+        resourceId,
+        workflowName,
+      });
+      expect(Array.isArray(runs)).toBe(true);
+      expect(runs.length).toBeGreaterThanOrEqual(2);
+      for (const run of runs) {
+        expect(run.resourceId).toBe(resourceId);
+      }
+    });
+
+    it('should return an empty array if no workflow runs match resourceId', async () => {
+      const { runs } = await store.getWorkflowRuns({
+        resourceId: 'non-existent-resource',
+        workflowName,
+      });
+      expect(Array.isArray(runs)).toBe(true);
+      expect(runs.length).toBe(0);
     });
   });
 
@@ -699,6 +838,38 @@ describe('PostgresStore', () => {
     });
   });
 
+  describe('hasColumn', () => {
+    const tempTable = 'temp_test_table';
+
+    beforeEach(async () => {
+      // Always try to drop the table before each test, ignore errors if it doesn't exist
+      try {
+        await store['db'].query(`DROP TABLE IF EXISTS ${tempTable}`);
+      } catch {
+        /* ignore */
+      }
+    });
+
+    it('returns true if the column exists', async () => {
+      await store['db'].query(`CREATE TABLE ${tempTable} (id SERIAL PRIMARY KEY, resourceId TEXT)`);
+      expect(await store['hasColumn'](tempTable, 'resourceId')).toBe(true);
+    });
+
+    it('returns false if the column does not exist', async () => {
+      await store['db'].query(`CREATE TABLE ${tempTable} (id SERIAL PRIMARY KEY)`);
+      expect(await store['hasColumn'](tempTable, 'resourceId')).toBe(false);
+    });
+
+    afterEach(async () => {
+      // Always try to drop the table after each test, ignore errors if it doesn't exist
+      try {
+        await store['db'].query(`DROP TABLE IF EXISTS ${tempTable}`);
+      } catch {
+        /* ignore */
+      }
+    });
+  });
+
   describe('Schema Support', () => {
     const customSchema = 'mastra_test';
     let customSchemaStore: PostgresStore;