Merge pull request #249 from launchql/feat/auth-method

pyramation · web-flow · commit a56adfab5ff6 · 2025-10-17T18:12:40.000-07:00
Feat/auth method
diff --git a/packages/pgsql-test/__tests__/postgres-test.auth-methods.test.ts b/packages/pgsql-test/__tests__/postgres-test.auth-methods.test.ts
@@ -0,0 +1,235 @@
+process.env.LOG_SCOPE = 'pgsql-test';
+
+import { getConnections } from '../src/connect';
+import { getRoleName } from '../src/roles';
+import { PgTestClient } from '../src/test-client';
+
+let db: PgTestClient;
+let pg: PgTestClient;
+let teardown: () => Promise<void>;
+
+beforeAll(async () => {
+  ({ db, pg, teardown } = await getConnections());
+  
+  await pg.query(`
+    CREATE TABLE IF NOT EXISTS test_users (
+      id SERIAL PRIMARY KEY,
+      email TEXT NOT NULL,
+      name TEXT
+    )
+  `);
+  
+  // Grant permissions to authenticated role
+  await pg.query(`GRANT ALL ON test_users TO authenticated`);
+  await pg.query(`GRANT ALL ON SEQUENCE test_users_id_seq TO authenticated`);
+});
+
+afterAll(async () => {
+  await pg.query(`DROP TABLE IF EXISTS test_users`);
+  await teardown();
+});
+
+describe('auth() method', () => {
+  beforeEach(async () => {
+    await db.beforeEach();
+  });
+
+  afterEach(async () => {
+    await db.afterEach();
+  });
+
+  it('sets role and userId', async () => {
+    const authRole = getRoleName('authenticated');
+    db.auth({ role: authRole, userId: '12345' });
+
+    const role = await db.query('SELECT current_setting(\'role\', true) AS role');
+    expect(role.rows[0].role).toBe(authRole);
+
+    const userId = await db.query('SELECT current_setting(\'jwt.claims.user_id\', true) AS user_id');
+    expect(userId.rows[0].user_id).toBe('12345');
+  });
+
+  it('sets role with custom userIdKey', async () => {
+    const authRole = getRoleName('authenticated');
+    db.auth({ 
+      role: authRole, 
+      userId: 'custom-123',
+      userIdKey: 'app.user.id'
+    });
+
+    const role = await db.query('SELECT current_setting(\'role\', true) AS role');
+    expect(role.rows[0].role).toBe(authRole);
+
+    const userId = await db.query('SELECT current_setting(\'app.user.id\', true) AS user_id');
+    expect(userId.rows[0].user_id).toBe('custom-123');
+  });
+
+  it('handles numeric userId', async () => {
+    const authRole = getRoleName('authenticated');
+    db.auth({ role: authRole, userId: 99999 });
+
+    const userId = await db.query('SELECT current_setting(\'jwt.claims.user_id\', true) AS user_id');
+    expect(userId.rows[0].user_id).toBe('99999');
+  });
+
+  it('sets role without userId', async () => {
+    const authRole = getRoleName('authenticated');
+    db.auth({ role: authRole });
+
+    const role = await db.query('SELECT current_setting(\'role\', true) AS role');
+    expect(role.rows[0].role).toBe(authRole);
+  });
+});
+
+describe('auth() with default role', () => {
+  beforeEach(async () => {
+    await db.beforeEach();
+  });
+
+  afterEach(async () => {
+    await db.afterEach();
+  });
+
+  it('uses default authenticated role when role not provided', async () => {
+    db.auth({ userId: 'test-user-456' });
+
+    const role = await db.query('SELECT current_setting(\'role\', true) AS role');
+    const expectedRole = getRoleName('authenticated');
+    expect(role.rows[0].role).toBe(expectedRole);
+
+    const userId = await db.query('SELECT current_setting(\'jwt.claims.user_id\', true) AS user_id');
+    expect(userId.rows[0].user_id).toBe('test-user-456');
+  });
+
+  it('allows explicit role override', async () => {
+    const anonRole = getRoleName('anonymous');
+    db.auth({ userId: 'admin-user-789', role: anonRole });
+
+    const role = await db.query('SELECT current_setting(\'role\', true) AS role');
+    expect(role.rows[0].role).toBe(anonRole);
+
+    const userId = await db.query('SELECT current_setting(\'jwt.claims.user_id\', true) AS user_id');
+    expect(userId.rows[0].user_id).toBe('admin-user-789');
+  });
+
+  it('handles numeric userId with default role', async () => {
+    db.auth({ userId: 42 });
+
+    const userId = await db.query('SELECT current_setting(\'jwt.claims.user_id\', true) AS user_id');
+    expect(userId.rows[0].user_id).toBe('42');
+  });
+});
+
+describe('clearContext() method', () => {
+  beforeEach(async () => {
+    await db.beforeEach();
+  });
+
+  afterEach(async () => {
+    await db.afterEach();
+  });
+
+  it('clears all session variables', async () => {
+    const authRole = getRoleName('authenticated');
+    db.setContext({
+      role: authRole,
+      'jwt.claims.user_id': 'test-123',
+      'jwt.claims.ip_address': '192.168.1.1',
+      'custom.var': 'custom-value'
+    });
+
+    const roleBefore = await db.query('SELECT current_setting(\'role\', true) AS role');
+    expect(roleBefore.rows[0].role).toBe(authRole);
+
+    const userIdBefore = await db.query('SELECT current_setting(\'jwt.claims.user_id\', true) AS user_id');
+    expect(userIdBefore.rows[0].user_id).toBe('test-123');
+
+    db.clearContext();
+
+    const defaultRole = getRoleName('anonymous');
+    const roleAfter = await db.query('SELECT current_setting(\'role\', true) AS role');
+    expect(roleAfter.rows[0].role).toBe(defaultRole);
+
+    const userIdAfter = await db.query('SELECT current_setting(\'jwt.claims.user_id\', true) AS user_id');
+    expect(userIdAfter.rows[0].user_id).toBe('');
+
+  });
+
+  it('allows setting new context after clearing', async () => {
+    const authRole = getRoleName('authenticated');
+    db.setContext({
+      role: authRole,
+      'jwt.claims.user_id': 'old-user'
+    });
+
+    const userIdBefore = await db.query('SELECT current_setting(\'jwt.claims.user_id\', true) AS user_id');
+    expect(userIdBefore.rows[0].user_id).toBe('old-user');
+
+    db.clearContext();
+    db.auth({ userId: 'new-user' });
+
+    const userId = await db.query('SELECT current_setting(\'jwt.claims.user_id\', true) AS user_id');
+    expect(userId.rows[0].user_id).toBe('new-user');
+  });
+});
+
+describe('publish() method', () => {
+  beforeEach(async () => {
+    await pg.beforeEach();
+    await db.beforeEach();
+  });
+
+  afterEach(async () => {
+    await db.afterEach();
+    await pg.afterEach();
+  });
+
+  it('makes data visible to other connections after publish', async () => {
+    await pg.query(`INSERT INTO test_users (email, name) VALUES ('alice@test.com', 'Alice')`);
+
+    const beforePublish = await db.any('SELECT * FROM test_users WHERE email = $1', ['alice@test.com']);
+    expect(beforePublish.length).toBe(0);
+
+    await pg.publish();
+
+    const afterPublish = await db.any('SELECT * FROM test_users WHERE email = $1', ['alice@test.com']);
+    expect(afterPublish.length).toBe(1);
+    expect(afterPublish[0].email).toBe('alice@test.com');
+    expect(afterPublish[0].name).toBe('Alice');
+  });
+
+  it('preserves context after publish', async () => {
+    const authRole = getRoleName('authenticated');
+    db.auth({ role: authRole, userId: 'test-999' });
+
+    await db.publish();
+
+    const role = await db.query('SELECT current_setting(\'role\', true) AS role');
+    expect(role.rows[0].role).toBe(authRole);
+
+    const userId = await db.query('SELECT current_setting(\'jwt.claims.user_id\', true) AS user_id');
+    expect(userId.rows[0].user_id).toBe('test-999');
+  });
+
+  it('allows rollback after publish', async () => {
+    await pg.query(`INSERT INTO test_users (email, name) VALUES ('bob@test.com', 'Bob')`);
+    await pg.publish();  // Bob is now committed and cannot be rolled back (important-comment)
+
+    await pg.query(`INSERT INTO test_users (email, name) VALUES ('charlie@test.com', 'Charlie')`);
+
+    const beforeRollback = await pg.any('SELECT * FROM test_users WHERE email IN ($1, $2)', ['bob@test.com', 'charlie@test.com']);
+    expect(beforeRollback.length).toBe(2);
+
+    await pg.rollback();  // Rollback Charlie to the savepoint created by publish() (important-comment)
+
+    const afterRollback = await pg.any('SELECT * FROM test_users WHERE email IN ($1, $2)', ['bob@test.com', 'charlie@test.com']);
+    expect(afterRollback.length).toBe(1);  // Only Bob remains (was committed via publish) (important-comment)
+    expect(afterRollback[0].email).toBe('bob@test.com');
+    
+    // Clean up Bob so it doesn't leak to other tests
+    await pg.query(`DELETE FROM test_users WHERE email = 'bob@test.com'`);
+    await pg.commit();  // Commit the deletion (important-comment)
+    await pg.begin();  // Start new transaction for outer afterEach (important-comment)
+    await pg.savepoint();  // Create savepoint for outer afterEach (important-comment)
+  });
+});
diff --git a/packages/pgsql-test/src/connect.ts b/packages/pgsql-test/src/connect.ts
@@ -102,10 +102,15 @@ export const getConnections = async (
     }
   }
 
-  const db = manager.getClient({
+  const dbConfig = {
     ...config,
     user: connOpts.connection.user,
     password: connOpts.connection.password
+  } as PgConfig;
+  
+  const db = manager.getClient(dbConfig, {
+    auth: connOpts.auth,
+    roles: connOpts.roles
   });
   db.setContext({ role: getDefaultRole(connOpts) });
 
diff --git a/packages/pgsql-test/src/manager.ts b/packages/pgsql-test/src/manager.ts
@@ -3,7 +3,7 @@ import { Pool } from 'pg';
 import { getPgEnvOptions, PgConfig } from 'pg-env';
 
 import { DbAdmin } from './admin';
-import { PgTestClient } from './test-client';
+import { PgTestClient, PgTestClientOpts } from './test-client';
 
 const log = new Logger('test-connector');
 
@@ -84,11 +84,14 @@ export class PgTestConnector {
     return this.pgPools.get(key)!;
   }
 
-  getClient(config: PgConfig): PgTestClient {
+  getClient(config: PgConfig, opts: Partial<PgTestClientOpts> = {}): PgTestClient {
     if (this.shuttingDown) {
       throw new Error('PgTestConnector is shutting down; no new clients allowed');
     }
-    const client = new PgTestClient(config, { trackConnect: (p) => this.registerConnect(p) });
+    const client = new PgTestClient(config, { 
+      trackConnect: (p) => this.registerConnect(p),
+      ...opts
+    });
     this.clients.add(client);
 
     const key = this.dbKey(config);
diff --git a/packages/pgsql-test/src/test-client.ts b/packages/pgsql-test/src/test-client.ts
@@ -1,19 +1,24 @@
 import { Client, QueryResult } from 'pg';
 import { PgConfig } from 'pg-env';
+import { AuthOptions, PgTestConnectionOptions } from '@launchql/types';
+import { getRoleName } from './roles';
 
-type PgTestClientOpts = {
+export type PgTestClientOpts = {
   deferConnect?: boolean;
   trackConnect?: (p: Promise<any>) => void;
-};
+} & Partial<PgTestConnectionOptions>;
 
 export class PgTestClient {
   public config: PgConfig;
   public client: Client;
+  private opts: PgTestClientOpts;
   private ctxStmts: string = '';
+  private contextSettings: Record<string, string | null> = {};
   private _ended: boolean = false;
   private connectPromise: Promise<void> | null = null;
 
   constructor(config: PgConfig, opts: PgTestClientOpts = {}) {
+    this.opts = opts;
     this.config = config;
     this.client = new Client({
       host: this.config.host,
@@ -71,7 +76,9 @@ export class PgTestClient {
   }
 
   setContext(ctx: Record<string, string | null>): void {
-    this.ctxStmts = Object.entries(ctx)
+    Object.assign(this.contextSettings, ctx);
+    
+    this.ctxStmts = Object.entries(this.contextSettings)
       .map(([key, val]) =>
         val === null
           ? `SELECT set_config('${key}', NULL, true);`
@@ -80,6 +87,59 @@ export class PgTestClient {
       .join('\n');
   }
 
+  /**
+   * Set authentication context for the current session.
+   * Configures role and user ID using cascading defaults from options → opts.auth → RoleMapping.
+   */
+  auth(options: AuthOptions = {}): void {
+    const role =
+      options.role ?? this.opts.auth?.role ?? getRoleName('authenticated', this.opts);
+    const userIdKey =
+      options.userIdKey ?? this.opts.auth?.userIdKey ?? 'jwt.claims.user_id';
+    const userId =
+      options.userId ?? this.opts.auth?.userId ?? null;
+
+    this.setContext({
+      role,
+      [userIdKey]: userId !== null ? String(userId) : null
+    });
+  }
+
+  /**
+   * Commit current transaction to make data visible to other connections, then start fresh transaction.
+   * Maintains test isolation by creating a savepoint and reapplying session context.
+   */
+  async publish(): Promise<void> {
+    await this.commit();    // make data visible to other sessions
+    await this.begin();     // fresh tx
+    await this.savepoint(); // keep rollback harness
+    await this.ctxQuery();  // reapply all setContext()
+  }
+
+  /**
+   * Clear all session context variables and reset to default anonymous role.
+   */
+  clearContext(): void {
+    const defaultRole = getRoleName('anonymous', this.opts);
+    
+    const nulledSettings: Record<string, string | null> = {};
+    Object.keys(this.contextSettings).forEach(key => {
+      nulledSettings[key] = null;
+    });
+    
+    nulledSettings.role = defaultRole;
+    
+    this.ctxStmts = Object.entries(nulledSettings)
+      .map(([key, val]) =>
+        val === null
+          ? `SELECT set_config('${key}', NULL, true);`
+          : `SELECT set_config('${key}', '${val}', true);`
+      )
+      .join('\n');
+    
+    this.contextSettings = { role: defaultRole };
+  }
+
   async any<T = any>(query: string, values?: any[]): Promise<T[]> {
     const result = await this.query(query, values);
     return result.rows;
diff --git a/packages/types/src/launchql.ts b/packages/types/src/launchql.ts
