@mastra/pg 0.12.3 → 0.12.4-alpha.0

package/src/storage/test-utils.ts

@@ -0,0 +1,368 @@
+import { createSampleThread } from '@internal/storage-test-utils';
+import type { StorageColumn, TABLE_NAMES } from '@mastra/core/storage';
+import pgPromise from 'pg-promise';
+import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it } from 'vitest';
+import { PostgresStore } from '.';
+import type { PostgresConfig } from '.';
+
+export const TEST_CONFIG: PostgresConfig = {
+  host: process.env.POSTGRES_HOST || 'localhost',
+  port: Number(process.env.POSTGRES_PORT) || 5434,
+  database: process.env.POSTGRES_DB || 'postgres',
+  user: process.env.POSTGRES_USER || 'postgres',
+  password: process.env.POSTGRES_PASSWORD || 'postgres',
+};
+
+export const connectionString = `postgresql://${TEST_CONFIG.user}:${TEST_CONFIG.password}@${TEST_CONFIG.host}:${TEST_CONFIG.port}/${TEST_CONFIG.database}`;
+
+export function pgTests() {
+  let store: PostgresStore;
+
+  describe('PG specific tests', () => {
+    beforeAll(async () => {
+      store = new PostgresStore(TEST_CONFIG);
+    });
+    afterAll(async () => {
+      try {
+        await store.close();
+      } catch {}
+    });
+
+    describe('Public Fields Access', () => {
+      it('should expose db field as public', () => {
+        expect(store.db).toBeDefined();
+        expect(typeof store.db).toBe('object');
+        expect(store.db.query).toBeDefined();
+        expect(typeof store.db.query).toBe('function');
+      });
+
+      it('should expose pgp field as public', () => {
+        expect(store.pgp).toBeDefined();
+        expect(typeof store.pgp).toBe('function');
+        expect(store.pgp.end).toBeDefined();
+        expect(typeof store.pgp.end).toBe('function');
+      });
+
+      it('should allow direct database queries via public db field', async () => {
+        const result = await store.db.one('SELECT 1 as test');
+        expect(result.test).toBe(1);
+      });
+
+      it('should allow access to pgp utilities via public pgp field', () => {
+        const helpers = store.pgp.helpers;
+        expect(helpers).toBeDefined();
+        expect(helpers.insert).toBeDefined();
+        expect(helpers.update).toBeDefined();
+      });
+
+      it('should maintain connection state through public db field', async () => {
+        // Test multiple queries to ensure connection state
+        const result1 = await store.db.one('SELECT NOW() as timestamp1');
+        const result2 = await store.db.one('SELECT NOW() as timestamp2');
+
+        expect(result1.timestamp1).toBeDefined();
+        expect(result2.timestamp2).toBeDefined();
+        expect(new Date(result2.timestamp2).getTime()).toBeGreaterThanOrEqual(new Date(result1.timestamp1).getTime());
+      });
+
+      it('should throw error when pool is used after disconnect', async () => {
+        await store.close();
+        await expect(store.db.connect()).rejects.toThrow();
+        store = new PostgresStore(TEST_CONFIG);
+        await store.init();
+      });
+    });
+
+    describe('PgStorage Table Name Quoting', () => {
+      const camelCaseTable = 'TestCamelCaseTable';
+      const snakeCaseTable = 'test_snake_case_table';
+      const BASE_SCHEMA = {
+        id: { type: 'integer', primaryKey: true, nullable: false },
+        name: { type: 'text', nullable: true },
+      } as Record<string, StorageColumn>;
+
+      beforeEach(async () => {
+        // Only clear tables if store is initialized
+        try {
+          // Clear tables before each test
+          await store.clearTable({ tableName: camelCaseTable as TABLE_NAMES });
+          await store.clearTable({ tableName: snakeCaseTable as TABLE_NAMES });
+        } catch (error) {
+          // Ignore errors during table clearing
+          console.warn('Error clearing tables:', error);
+        }
+      });
+
+      afterEach(async () => {
+        // Only clear tables if store is initialized
+        try {
+          // Clear tables before each test
+          await store.clearTable({ tableName: camelCaseTable as TABLE_NAMES });
+          await store.clearTable({ tableName: snakeCaseTable as TABLE_NAMES });
+        } catch (error) {
+          // Ignore errors during table clearing
+          console.warn('Error clearing tables:', error);
+        }
+      });
+
+      it('should create and upsert to a camelCase table without quoting errors', async () => {
+        await expect(
+          store.createTable({
+            tableName: camelCaseTable as TABLE_NAMES,
+            schema: BASE_SCHEMA,
+          }),
+        ).resolves.not.toThrow();
+
+        await store.insert({
+          tableName: camelCaseTable as TABLE_NAMES,
+          record: { id: '1', name: 'Alice' },
+        });
+
+        const row: any = await store.load({
+          tableName: camelCaseTable as TABLE_NAMES,
+          keys: { id: '1' },
+        });
+        expect(row?.name).toBe('Alice');
+      });
+
+      it('should create and upsert to a snake_case table without quoting errors', async () => {
+        await expect(
+          store.createTable({
+            tableName: snakeCaseTable as TABLE_NAMES,
+            schema: BASE_SCHEMA,
+          }),
+        ).resolves.not.toThrow();
+
+        await store.insert({
+          tableName: snakeCaseTable as TABLE_NAMES,
+          record: { id: '2', name: 'Bob' },
+        });
+
+        const row: any = await store.load({
+          tableName: snakeCaseTable as TABLE_NAMES,
+          keys: { id: '2' },
+        });
+        expect(row?.name).toBe('Bob');
+      });
+    });
+
+    describe('Permission Handling', () => {
+      const schemaRestrictedUser = 'mastra_schema_restricted_storage';
+      const restrictedPassword = 'test123';
+      const testSchema = 'testSchema';
+      let adminDb: pgPromise.IDatabase<{}>;
+      let pgpAdmin: pgPromise.IMain;
+
+      beforeAll(async () => {
+        // Re-initialize the main store for subsequent tests
+
+        await store.init();
+
+        // 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 as any).database} FROM ${schemaRestrictedUser};
+                  GRANT CONNECT ON DATABASE ${(TEST_CONFIG as any).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 {
+          // 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);
+          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,
+            schemaName: 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,
+            schemaName: 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
+          }
+        });
+      });
+    });
+
+    describe('Validation', () => {
+      const validConfig = TEST_CONFIG as any;
+      it('throws if connectionString is empty', () => {
+        expect(() => new PostgresStore({ connectionString: '' })).toThrow();
+        expect(() => new PostgresStore({ ...validConfig, connectionString: '' })).toThrow();
+      });
+      it('throws if host is missing or empty', () => {
+        expect(() => new PostgresStore({ ...validConfig, host: '' })).toThrow();
+        const { host, ...rest } = validConfig;
+        expect(() => new PostgresStore(rest as any)).toThrow();
+      });
+      it('throws if connectionString is empty', () => {
+        expect(() => new PostgresStore({ connectionString: '' })).toThrow();
+        const { database, ...rest } = validConfig;
+        expect(() => new PostgresStore(rest as any)).toThrow();
+      });
+      it('throws if user is missing or empty', () => {
+        expect(() => new PostgresStore({ ...validConfig, user: '' })).toThrow();
+        const { user, ...rest } = validConfig;
+        expect(() => new PostgresStore(rest as any)).toThrow();
+      });
+      it('throws if database is missing or empty', () => {
+        expect(() => new PostgresStore({ ...validConfig, database: '' })).toThrow();
+        const { database, ...rest } = validConfig;
+        expect(() => new PostgresStore(rest as any)).toThrow();
+      });
+      it('throws if password is missing or empty', () => {
+        expect(() => new PostgresStore({ ...validConfig, password: '' })).toThrow();
+        const { password, ...rest } = validConfig;
+        expect(() => new PostgresStore(rest as any)).toThrow();
+      });
+      it('does not throw on valid config (host-based)', () => {
+        expect(() => new PostgresStore(validConfig)).not.toThrow();
+      });
+      it('does not throw on non-empty connection string', () => {
+        expect(() => new PostgresStore({ connectionString })).not.toThrow();
+      });
+    });
+  });
+}

package/src/vector/index.test.ts

@@ -1988,6 +1988,95 @@ describe('PgVector', () => {
         expect(results[0]?.score).toBeCloseTo(1, 5);
         expect(results[1]?.score).toBeGreaterThan(0.9);
       });
+
+      // NEW TEST: Reproduce the SET LOCAL bug
+      it('should verify that ef_search parameter is actually being set (reproduces SET LOCAL bug)', async () => {
+        const client = await vectorDB.pool.connect();
+        try {
+          // Test current behavior: SET LOCAL without transaction should have no effect
+          await client.query('SET LOCAL hnsw.ef_search = 500');
+
+          // Check if the parameter was actually set
+          const result = await client.query('SHOW hnsw.ef_search');
+          const currentValue = result.rows[0]['hnsw.ef_search'];
+
+          // The value should still be the default (not 500)
+          expect(parseInt(currentValue)).not.toBe(500);
+
+          // Now test with proper transaction
+          await client.query('BEGIN');
+          await client.query('SET LOCAL hnsw.ef_search = 500');
+
+          const resultInTransaction = await client.query('SHOW hnsw.ef_search');
+          const valueInTransaction = resultInTransaction.rows[0]['hnsw.ef_search'];
+
+          // This should work because we're in a transaction
+          expect(parseInt(valueInTransaction)).toBe(500);
+
+          await client.query('ROLLBACK');
+
+          // After rollback, should return to default
+          const resultAfterRollback = await client.query('SHOW hnsw.ef_search');
+          const valueAfterRollback = resultAfterRollback.rows[0]['hnsw.ef_search'];
+          expect(parseInt(valueAfterRollback)).not.toBe(500);
+        } finally {
+          client.release();
+        }
+      });
+
+      // Verify the fix works - ef parameter is properly applied in query method
+      it('should properly apply ef parameter using transactions (verifies fix)', async () => {
+        const client = await vectorDB.pool.connect();
+        const queryCommands: string[] = [];
+
+        // Spy on the client query method to capture all SQL commands
+        const originalClientQuery = client.query;
+        const clientQuerySpy = vi.fn().mockImplementation((query, ...args) => {
+          if (typeof query === 'string') {
+            queryCommands.push(query);
+          }
+          return originalClientQuery.call(client, query, ...args);
+        });
+        client.query = clientQuerySpy;
+
+        try {
+          // Manually release the client so query() can get a fresh one
+          client.release();
+
+          await vectorDB.query({
+            indexName,
+            queryVector: [1, 0, 0],
+            topK: 2,
+            ef: 128,
+          });
+
+          const testClient = await vectorDB.pool.connect();
+          try {
+            // Test that SET LOCAL works within a transaction
+            await testClient.query('BEGIN');
+            await testClient.query('SET LOCAL hnsw.ef_search = 256');
+
+            const result = await testClient.query('SHOW hnsw.ef_search');
+            const value = result.rows[0]['hnsw.ef_search'];
+            expect(parseInt(value)).toBe(256);
+
+            await testClient.query('ROLLBACK');
+
+            // After rollback, should revert
+            const resultAfter = await testClient.query('SHOW hnsw.ef_search');
+            const valueAfter = resultAfter.rows[0]['hnsw.ef_search'];
+            expect(parseInt(valueAfter)).not.toBe(256);
+          } finally {
+            testClient.release();
+          }
+        } finally {
+          // Restore original function if client is still connected
+          if (client.query === clientQuerySpy) {
+            client.query = originalClientQuery;
+          }
+          clientQuerySpy.mockRestore();
+        }
+      });
     });
 
     describe('IVF Parameters', () => {

package/src/vector/index.ts

@@ -200,6 +200,7 @@ export class PgVector extends MastraVector<PGVectorFilter> {
 
     const client = await this.pool.connect();
     try {
+      await client.query('BEGIN');
       const vectorStr = `[${queryVector.join(',')}]`;
       const translatedFilter = this.transformFilter(filter);
       const { sql: filterQuery, values: filterValues } = buildFilterQuery(translatedFilter, minScore, topK);
@@ -237,6 +238,7 @@ export class PgVector extends MastraVector<PGVectorFilter> {
         ORDER BY score DESC
         LIMIT $2`;
       const result = await client.query(query, filterValues);
+      await client.query('COMMIT');
 
       return result.rows.map(({ id, score, metadata, embedding }) => ({
         id,
@@ -245,6 +247,7 @@ export class PgVector extends MastraVector<PGVectorFilter> {
         ...(includeVector && embedding && { vector: JSON.parse(embedding) }),
       }));
     } catch (error) {
+      await client.query('ROLLBACK');
       const mastraError = new MastraError(
         {
           id: 'MASTRA_STORAGE_PG_VECTOR_QUERY_FAILED',