@mastra/pg 0.2.10-alpha.3 → 0.2.10-alpha.5

package/src/storage/index.test.ts

@@ -1,8 +1,10 @@
 import { randomUUID } from 'crypto';
 import type { MetricResult } from '@mastra/core/eval';
+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 { describe, it, expect, beforeAll, beforeEach, afterAll } from 'vitest';
+import pgPromise from 'pg-promise';
+import { describe, it, expect, beforeAll, beforeEach, afterAll, afterEach } from 'vitest';
 
 import { PostgresStore } from '.';
 import type { PostgresConfig } from '.';
@@ -15,6 +17,8 @@ const TEST_CONFIG: PostgresConfig = {
   password: process.env.POSTGRES_PASSWORD || 'postgres',
 };
 
+const connectionString = `postgresql://${TEST_CONFIG.user}:${TEST_CONFIG.password}@${TEST_CONFIG.host}:${TEST_CONFIG.port}/${TEST_CONFIG.database}`;
+
 // Sample test data factory functions
 const createSampleThread = () => ({
   id: `thread-${randomUUID()}`,
@@ -87,11 +91,17 @@ describe('PostgresStore', () => {
   });
 
   beforeEach(async () => {
-    // Clear tables before each test
-    await store.clearTable({ tableName: TABLE_WORKFLOW_SNAPSHOT });
-    await store.clearTable({ tableName: TABLE_MESSAGES });
-    await store.clearTable({ tableName: TABLE_THREADS });
-    await store.clearTable({ tableName: TABLE_EVALS });
+    // Only clear tables if store is initialized
+    try {
+      // Clear tables before each test
+      await store.clearTable({ tableName: TABLE_WORKFLOW_SNAPSHOT });
+      await store.clearTable({ tableName: TABLE_MESSAGES });
+      await store.clearTable({ tableName: TABLE_THREADS });
+      await store.clearTable({ tableName: TABLE_EVALS });
+    } catch (error) {
+      // Ignore errors during table clearing
+      console.warn('Error clearing tables:', error);
+    }
   });
 
   describe('Thread Operations', () => {
@@ -192,9 +202,9 @@ describe('PostgresStore', () => {
       await store.__saveThread({ thread });
 
       const messages = [
-        { ...createSampleMessage(thread.id), content: [{ type: 'text', text: 'First' }] },
-        { ...createSampleMessage(thread.id), content: [{ type: 'text', text: 'Second' }] },
-        { ...createSampleMessage(thread.id), content: [{ type: 'text', text: 'Third' }] },
+        { ...createSampleMessage(thread.id), content: [{ type: 'text', text: 'First' }] as MessageType['content'] },
+        { ...createSampleMessage(thread.id), content: [{ type: 'text', text: 'Second' }] as MessageType['content'] },
+        { ...createSampleMessage(thread.id), content: [{ type: 'text', text: 'Third' }] as MessageType['content'] },
       ];
 
       await store.__saveMessages({ messages });
@@ -204,7 +214,7 @@ describe('PostgresStore', () => {
 
       // Verify order is maintained
       retrievedMessages.forEach((msg, idx) => {
-        expect(msg.content[0].text).toBe(messages[idx].content[0].text);
+        expect((msg.content[0] as any).text).toBe((messages[idx].content[0] as any).text);
       });
     });
 
@@ -214,7 +224,7 @@ describe('PostgresStore', () => {
 
       const messages = [
         createSampleMessage(thread.id),
-        { ...createSampleMessage(thread.id), id: null }, // This will cause an error
+        { ...createSampleMessage(thread.id), id: null } as any, // This will cause an error
       ];
 
       await expect(store.__saveMessages({ messages })).rejects.toThrow();
@@ -645,7 +655,262 @@ describe('PostgresStore', () => {
     });
   });
 
+  describe('Schema Support', () => {
+    const customSchema = 'mastra_test';
+    let customSchemaStore: PostgresStore;
+
+    beforeAll(async () => {
+      customSchemaStore = new PostgresStore({
+        ...TEST_CONFIG,
+        schema: customSchema,
+      });
+
+      await customSchemaStore.init();
+    });
+
+    afterAll(async () => {
+      await customSchemaStore.close();
+      // Re-initialize the main store for subsequent tests
+      store = new PostgresStore(TEST_CONFIG);
+      await store.init();
+    });
+
+    describe('Constructor and Initialization', () => {
+      it('should accept connectionString directly', () => {
+        // Use existing store instead of creating new one
+        expect(store).toBeInstanceOf(PostgresStore);
+      });
+
+      it('should accept config object with schema', () => {
+        // Use existing custom schema store
+        expect(customSchemaStore).toBeInstanceOf(PostgresStore);
+      });
+    });
+
+    describe('Schema Operations', () => {
+      it('should create and query tables in custom schema', async () => {
+        // Create thread in custom schema
+        const thread = createSampleThread();
+        await customSchemaStore.__saveThread({ thread });
+
+        // Verify thread exists in custom schema
+        const retrieved = await customSchemaStore.__getThreadById({ threadId: thread.id });
+        expect(retrieved?.title).toBe(thread.title);
+      });
+
+      it('should allow same table names in different schemas', async () => {
+        // Create threads in both schemas
+        const defaultThread = createSampleThread();
+        const customThread = createSampleThread();
+
+        await store.__saveThread({ thread: defaultThread });
+        await customSchemaStore.__saveThread({ thread: customThread });
+
+        // Verify threads exist in respective schemas
+        const defaultResult = await store.__getThreadById({ threadId: defaultThread.id });
+        const customResult = await customSchemaStore.__getThreadById({ threadId: customThread.id });
+
+        expect(defaultResult?.id).toBe(defaultThread.id);
+        expect(customResult?.id).toBe(customThread.id);
+
+        // Verify cross-schema isolation
+        const defaultInCustom = await customSchemaStore.__getThreadById({ threadId: defaultThread.id });
+        const customInDefault = await store.__getThreadById({ threadId: customThread.id });
+
+        expect(defaultInCustom).toBeNull();
+        expect(customInDefault).toBeNull();
+      });
+    });
+  });
+
+  describe('Permission Handling', () => {
+    const schemaRestrictedUser = 'mastra_schema_restricted_storage';
+    const restrictedPassword = 'test123';
+    const testSchema = 'test_schema';
+    let adminDb: pgPromise.IDatabase<{}>;
+    let pgpAdmin: pgPromise.IMain;
+
+    beforeAll(async () => {
+      // Create a separate pg-promise instance for admin operations
+      pgpAdmin = pgPromise();
+      adminDb = pgpAdmin(connectionString);
+      try {
+        await adminDb.tx(async t => {
+          // Drop the test schema if it exists from previous runs
+          await t.none(`DROP SCHEMA IF EXISTS ${testSchema} CASCADE`);
+
+          // Create schema restricted user with minimal permissions
+          await t.none(`          
+          DO $$
+          BEGIN
+            IF NOT EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = '${schemaRestrictedUser}') THEN
+              CREATE USER ${schemaRestrictedUser} WITH PASSWORD '${restrictedPassword}' NOCREATEDB;
+            END IF;
+          END
+          $$;`);
+
+          // Grant only connect and usage to schema restricted user
+          await t.none(`
+            REVOKE ALL ON DATABASE ${TEST_CONFIG.database} FROM ${schemaRestrictedUser};
+            GRANT CONNECT ON DATABASE ${TEST_CONFIG.database} TO ${schemaRestrictedUser};
+            REVOKE ALL ON SCHEMA public FROM ${schemaRestrictedUser};
+            GRANT USAGE ON SCHEMA public TO ${schemaRestrictedUser};
+          `);
+        });
+      } catch (error) {
+        // Clean up the database connection on error
+        pgpAdmin.end();
+        throw error;
+      }
+    });
+
+    afterAll(async () => {
+      try {
+        // First close any store connections
+        if (store) {
+          await store.close();
+        }
+
+        // Then clean up test user in admin connection
+        await adminDb.tx(async t => {
+          await t.none(`
+            REASSIGN OWNED BY ${schemaRestrictedUser} TO postgres;
+            DROP OWNED BY ${schemaRestrictedUser};
+            DROP USER IF EXISTS ${schemaRestrictedUser};
+          `);
+        });
+
+        // Finally clean up admin connection
+        if (pgpAdmin) {
+          pgpAdmin.end();
+        }
+      } catch (error) {
+        console.error('Error cleaning up test user:', error);
+        // Still try to clean up connections even if user cleanup fails
+        if (store) await store.close();
+        if (pgpAdmin) pgpAdmin.end();
+      }
+    });
+
+    describe('Schema Creation', () => {
+      beforeEach(async () => {
+        // Create a fresh connection for each test
+        const tempPgp = pgPromise();
+        const tempDb = tempPgp(connectionString);
+
+        try {
+          // Ensure schema doesn't exist before each test
+          await tempDb.none(`DROP SCHEMA IF EXISTS ${testSchema} CASCADE`);
+
+          // Ensure no active connections from restricted user
+          await tempDb.none(`
+            SELECT pg_terminate_backend(pid) 
+            FROM pg_stat_activity 
+            WHERE usename = '${schemaRestrictedUser}'
+          `);
+        } finally {
+          tempPgp.end(); // Always clean up the connection
+        }
+      });
+
+      afterEach(async () => {
+        // Create a fresh connection for cleanup
+        const tempPgp = pgPromise();
+        const tempDb = tempPgp(connectionString);
+
+        try {
+          // Clean up any connections from the restricted user and drop schema
+          await tempDb.none(`
+            DO $$
+            BEGIN
+              -- Terminate connections
+              PERFORM pg_terminate_backend(pid) 
+              FROM pg_stat_activity 
+              WHERE usename = '${schemaRestrictedUser}';
+
+              -- Drop schema
+              DROP SCHEMA IF EXISTS ${testSchema} CASCADE;
+            END $$;
+          `);
+        } catch (error) {
+          console.error('Error in afterEach cleanup:', error);
+        } finally {
+          tempPgp.end(); // Always clean up the connection
+        }
+      });
+
+      it('should fail when user lacks CREATE privilege', async () => {
+        const restrictedDB = new PostgresStore({
+          ...TEST_CONFIG,
+          user: schemaRestrictedUser,
+          password: restrictedPassword,
+          schema: testSchema,
+        });
+
+        // Create a fresh connection for verification
+        const tempPgp = pgPromise();
+        const tempDb = tempPgp(connectionString);
+
+        try {
+          // Test schema creation by initializing the store
+          await expect(async () => {
+            await restrictedDB.init();
+          }).rejects.toThrow(
+            `Unable to create schema "${testSchema}". This requires CREATE privilege on the database.`,
+          );
+
+          // Verify schema was not created
+          const exists = await tempDb.oneOrNone(
+            `SELECT EXISTS (SELECT 1 FROM information_schema.schemata WHERE schema_name = $1)`,
+            [testSchema],
+          );
+          expect(exists?.exists).toBe(false);
+        } finally {
+          await restrictedDB.close();
+          tempPgp.end(); // Clean up the verification connection
+        }
+      });
+
+      it('should fail with schema creation error when saving thread', async () => {
+        const restrictedDB = new PostgresStore({
+          ...TEST_CONFIG,
+          user: schemaRestrictedUser,
+          password: restrictedPassword,
+          schema: testSchema,
+        });
+
+        // Create a fresh connection for verification
+        const tempPgp = pgPromise();
+        const tempDb = tempPgp(connectionString);
+
+        try {
+          await expect(async () => {
+            await restrictedDB.init();
+            const thread = createSampleThread();
+            await restrictedDB.__saveThread({ thread });
+          }).rejects.toThrow(
+            `Unable to create schema "${testSchema}". This requires CREATE privilege on the database.`,
+          );
+
+          // Verify schema was not created
+          const exists = await tempDb.oneOrNone(
+            `SELECT EXISTS (SELECT 1 FROM information_schema.schemata WHERE schema_name = $1)`,
+            [testSchema],
+          );
+          expect(exists?.exists).toBe(false);
+        } finally {
+          await restrictedDB.close();
+          tempPgp.end(); // Clean up the verification connection
+        }
+      });
+    });
+  });
+
   afterAll(async () => {
-    await store.close();
+    try {
+      await store.close();
+    } catch (error) {
+      console.warn('Error closing store:', error);
+    }
   });
 });

package/src/storage/index.ts

@@ -13,7 +13,7 @@ import type { WorkflowRunState } from '@mastra/core/workflows';
 import pgPromise from 'pg-promise';
 import type { ISSLConfig } from 'pg-promise/typescript/pg-subset';
 
-export type PostgresConfig =
+export type PostgresConfig = { schema?: string } & (
   | {
       host: string;
       port: number;
@@ -24,15 +24,20 @@ export type PostgresConfig =
     }
   | {
       connectionString: string;
-    };
+    }
+);
 
 export class PostgresStore extends MastraStorage {
   private db: pgPromise.IDatabase<{}>;
   private pgp: pgPromise.IMain;
+  private schema?: string;
+  private setupSchemaPromise: Promise<void> | null = null;
+  private schemaSetupComplete: boolean | undefined = undefined;
 
   constructor(config: PostgresConfig) {
     super({ name: 'PostgresStore' });
     this.pgp = pgPromise();
+    this.schema = config.schema;
     this.db = this.pgp(
       `connectionString` in config
         ? { connectionString: config.connectionString }
@@ -47,9 +52,13 @@ export class PostgresStore extends MastraStorage {
     );
   }
 
+  private getTableName(indexName: string) {
+    return this.schema ? `${this.schema}."${indexName}"` : `"${indexName}"`;
+  }
+
   async getEvalsByAgentName(agentName: string, type?: 'test' | 'live'): Promise<EvalRow[]> {
     try {
-      const baseQuery = `SELECT * FROM ${TABLE_EVALS} WHERE agent_name = $1`;
+      const baseQuery = `SELECT * FROM ${this.getTableName(TABLE_EVALS)} WHERE agent_name = $1`;
       const typeCondition =
         type === 'test'
           ? " AND test_info IS NOT NULL AND test_info->>'testPath' IS NOT NULL"
@@ -186,7 +195,10 @@ export class PostgresStore extends MastraStorage {
       endTime: string;
       other: any;
       createdAt: string;
-    }>(`SELECT * FROM ${TABLE_TRACES} ${whereClause} ORDER BY "createdAt" DESC LIMIT ${limit} OFFSET ${offset}`, args);
+    }>(
+      `SELECT * FROM ${this.getTableName(TABLE_TRACES)} ${whereClause} ORDER BY "createdAt" DESC LIMIT ${limit} OFFSET ${offset}`,
+      args,
+    );
 
     if (!result) {
       return [];
@@ -210,6 +222,55 @@ export class PostgresStore extends MastraStorage {
     })) as any;
   }
 
+  private async setupSchema() {
+    if (!this.schema || this.schemaSetupComplete) {
+      return;
+    }
+
+    if (!this.setupSchemaPromise) {
+      this.setupSchemaPromise = (async () => {
+        try {
+          // First check if schema exists and we have usage permission
+          const schemaExists = await this.db.oneOrNone(
+            `
+            SELECT EXISTS (
+              SELECT 1 FROM information_schema.schemata 
+              WHERE schema_name = $1
+            )
+          `,
+            [this.schema],
+          );
+
+          if (!schemaExists?.exists) {
+            try {
+              await this.db.none(`CREATE SCHEMA IF NOT EXISTS ${this.schema}`);
+              this.logger.info(`Schema "${this.schema}" created successfully`);
+            } catch (error) {
+              this.logger.error(`Failed to create schema "${this.schema}"`, { error });
+              throw new Error(
+                `Unable to create schema "${this.schema}". This requires CREATE privilege on the database. ` +
+                  `Either create the schema manually or grant CREATE privilege to the user.`,
+              );
+            }
+          }
+
+          // If we got here, schema exists and we can use it
+          this.schemaSetupComplete = true;
+          this.logger.debug(`Schema "${this.schema}" is ready for use`);
+        } catch (error) {
+          // Reset flags so we can retry
+          this.schemaSetupComplete = undefined;
+          this.setupSchemaPromise = null;
+          throw error;
+        } finally {
+          this.setupSchemaPromise = null;
+        }
+      })();
+    }
+
+    await this.setupSchemaPromise;
+  }
+
   async createTable({
     tableName,
     schema,
@@ -227,8 +288,13 @@ export class PostgresStore extends MastraStorage {
         })
         .join(',\n');
 
+      // Create schema if it doesn't exist
+      if (this.schema) {
+        await this.setupSchema();
+      }
+
       const sql = `
-        CREATE TABLE IF NOT EXISTS ${tableName} (
+        CREATE TABLE IF NOT EXISTS ${this.getTableName(tableName)} (
           ${columns}
         );
         ${
@@ -238,7 +304,7 @@ export class PostgresStore extends MastraStorage {
           IF NOT EXISTS (
             SELECT 1 FROM pg_constraint WHERE conname = 'mastra_workflow_snapshot_workflow_name_run_id_key'
           ) THEN
-            ALTER TABLE ${tableName}
+            ALTER TABLE ${this.getTableName(tableName)}
             ADD CONSTRAINT mastra_workflow_snapshot_workflow_name_run_id_key
             UNIQUE (workflow_name, run_id);
           END IF;
@@ -257,7 +323,7 @@ export class PostgresStore extends MastraStorage {
 
   async clearTable({ tableName }: { tableName: TABLE_NAMES }): Promise<void> {
     try {
-      await this.db.none(`TRUNCATE TABLE ${tableName} CASCADE`);
+      await this.db.none(`TRUNCATE TABLE ${this.getTableName(tableName)} CASCADE`);
     } catch (error) {
       console.error(`Error clearing table ${tableName}:`, error);
       throw error;
@@ -271,7 +337,7 @@ export class PostgresStore extends MastraStorage {
       const placeholders = values.map((_, i) => `$${i + 1}`).join(', ');
 
       await this.db.none(
-        `INSERT INTO ${tableName} (${columns.map(c => `"${c}"`).join(', ')}) VALUES (${placeholders})`,
+        `INSERT INTO ${this.getTableName(tableName)} (${columns.map(c => `"${c}"`).join(', ')}) VALUES (${placeholders})`,
         values,
       );
     } catch (error) {
@@ -286,7 +352,10 @@ export class PostgresStore extends MastraStorage {
       const conditions = keyEntries.map(([key], index) => `"${key}" = $${index + 1}`).join(' AND ');
       const values = keyEntries.map(([_, value]) => value);
 
-      const result = await this.db.oneOrNone<R>(`SELECT * FROM ${tableName} WHERE ${conditions}`, values);
+      const result = await this.db.oneOrNone<R>(
+        `SELECT * FROM ${this.getTableName(tableName)} WHERE ${conditions}`,
+        values,
+      );
 
       if (!result) {
         return null;
@@ -318,7 +387,7 @@ export class PostgresStore extends MastraStorage {
           metadata,
           "createdAt",
           "updatedAt"
-        FROM "${TABLE_THREADS}"
+        FROM ${this.getTableName(TABLE_THREADS)}
         WHERE id = $1`,
         [threadId],
       );
@@ -349,7 +418,7 @@ export class PostgresStore extends MastraStorage {
           metadata,
           "createdAt",
           "updatedAt"
-        FROM "${TABLE_THREADS}"
+        FROM ${this.getTableName(TABLE_THREADS)}
         WHERE "resourceId" = $1`,
         [resourceId],
       );
@@ -369,7 +438,7 @@ export class PostgresStore extends MastraStorage {
   async saveThread({ thread }: { thread: StorageThreadType }): Promise<StorageThreadType> {
     try {
       await this.db.none(
-        `INSERT INTO "${TABLE_THREADS}" (
+        `INSERT INTO ${this.getTableName(TABLE_THREADS)} (
           id,
           "resourceId",
           title,
@@ -423,7 +492,7 @@ export class PostgresStore extends MastraStorage {
       };
 
       const thread = await this.db.one<StorageThreadType>(
-        `UPDATE "${TABLE_THREADS}"
+        `UPDATE ${this.getTableName(TABLE_THREADS)}
         SET title = $1,
             metadata = $2,
             "updatedAt" = $3
@@ -448,10 +517,10 @@ export class PostgresStore extends MastraStorage {
     try {
       await this.db.tx(async t => {
         // First delete all messages associated with this thread
-        await t.none(`DELETE FROM "${TABLE_MESSAGES}" WHERE thread_id = $1`, [threadId]);
+        await t.none(`DELETE FROM ${this.getTableName(TABLE_MESSAGES)} WHERE thread_id = $1`, [threadId]);
 
         // Then delete the thread
-        await t.none(`DELETE FROM "${TABLE_THREADS}" WHERE id = $1`, [threadId]);
+        await t.none(`DELETE FROM ${this.getTableName(TABLE_THREADS)} WHERE id = $1`, [threadId]);
       });
     } catch (error) {
       console.error('Error deleting thread:', error);
@@ -472,7 +541,7 @@ export class PostgresStore extends MastraStorage {
             SELECT 
               *,
               ROW_NUMBER() OVER (ORDER BY "createdAt" DESC) as row_num
-            FROM "${TABLE_MESSAGES}"
+            FROM ${this.getTableName(TABLE_MESSAGES)}
             WHERE thread_id = $1
           )
           SELECT
@@ -518,7 +587,7 @@ export class PostgresStore extends MastraStorage {
             type,
             "createdAt", 
             thread_id AS "threadId"
-        FROM "${TABLE_MESSAGES}"
+        FROM ${this.getTableName(TABLE_MESSAGES)}
         WHERE thread_id = $1
         AND id != ALL($2)
         ORDER BY "createdAt" DESC
@@ -568,7 +637,7 @@ export class PostgresStore extends MastraStorage {
       await this.db.tx(async t => {
         for (const message of messages) {
           await t.none(
-            `INSERT INTO "${TABLE_MESSAGES}" (id, thread_id, content, "createdAt", role, type) 
+            `INSERT INTO ${this.getTableName(TABLE_MESSAGES)} (id, thread_id, content, "createdAt", role, type) 
              VALUES ($1, $2, $3, $4, $5, $6)`,
             [
               message.id,
@@ -601,7 +670,7 @@ export class PostgresStore extends MastraStorage {
     try {
       const now = new Date().toISOString();
       await this.db.none(
-        `INSERT INTO "${TABLE_WORKFLOW_SNAPSHOT}" (
+        `INSERT INTO ${this.getTableName(TABLE_WORKFLOW_SNAPSHOT)} (
           workflow_name,
           run_id,
           snapshot,
@@ -696,7 +765,7 @@ export class PostgresStore extends MastraStorage {
     // Only get total count when using pagination
     if (limit !== undefined && offset !== undefined) {
       const countResult = await this.db.one(
-        `SELECT COUNT(*) as count FROM ${TABLE_WORKFLOW_SNAPSHOT} ${whereClause}`,
+        `SELECT COUNT(*) as count FROM ${this.getTableName(TABLE_WORKFLOW_SNAPSHOT)} ${whereClause}`,
         values,
       );
       total = Number(countResult.count);
@@ -704,7 +773,7 @@ export class PostgresStore extends MastraStorage {
 
     // Get results
     const query = `
-      SELECT * FROM ${TABLE_WORKFLOW_SNAPSHOT} 
+      SELECT * FROM ${this.getTableName(TABLE_WORKFLOW_SNAPSHOT)} 
       ${whereClause} 
       ORDER BY "createdAt" DESC
       ${limit !== undefined && offset !== undefined ? ` LIMIT $${paramIndex} OFFSET $${paramIndex + 1}` : ''}