@accounter/server/src/modules/app-providers/tenant-db-client.ts

Accounter GraphQL server

import { AsyncLocalStorage } from 'node:async_hooks';
import { Mutex } from 'async-mutex';
import { GraphQLError } from 'graphql';
import { Injectable, Scope } from 'graphql-modules';
import type { PoolClient, QueryResult, QueryResultRow } from 'pg';
import { resolveWriteTargetBusinessId } from '../../shared/helpers/auth-scope.js';
import type { AuthContext } from '../../shared/types/auth.js';
import { AuthContextProvider } from '../auth/providers/auth-context.provider.js';
import { DBProvider } from './db.provider.js';

/**
 * TenantAwareDBClient enforces Row-Level Security (RLS) by setting PostgreSQL
 * session variables for every database transaction.
 
 * 
 * RLS Enforcement:
 * - app.current_business_id: Set to the authenticated user's active business
 * - app.current_user_id: Set to the authenticated user's ID (or NULL for API keys)
 * - app.auth_type: Set to 'jwt' or 'apiKey'
 *
 * **Usage:**
 * Inject into Operation-scoped providers via constructor DI:
 *
 * @example
 * @Injectable({ scope: Scope.Operation })
 * class BusinessesProvider {
 *   constructor(private db: TenantAwareDBClient) {}
 *
 *   async getBusinesses() {
 *     return this.db.query('SELECT * FROM businesses')
 *   }
 * }
 *
 * Transaction Management:
 * - Supports nested transactions via SAVEPOINTs
 * - Automatically rolls back on error
 * - Automatically releases connection on dispose
 *
 * **DO NOT** access from Yoga context - use DI injection instead.
 *
 * @throws {GraphQLError} UNAUTHENTICATED if auth context is null
 */
@Injectable({
  scope: Scope.Operation,
  global: true,
})
export class TenantAwareDBClient {
  private mutex = new Mutex();
  private storage = new AsyncLocalStorage<boolean>();
  private activeClient: PoolClient | null = null;
  private transactionDepth = 0;
  private isDisposed = false;
  private initializationPromise: Promise<void> | null = null;
  private authContext: AuthContext | null = null;
  private authContextInitialized = false;

  constructor(
    private dbProvider: DBProvider,
    private authContextProvider: AuthContextProvider,
  ) {}

  /**
   * Execute a query with RLS enforcement.
   * If a transaction is already active, uses it.
   * If not, starts a new transaction/session, executes the query, and commits.
   */
  public async query<T extends QueryResultRow = QueryResultRow>(
    text: string,
    params?: unknown[],
  ): Promise<QueryResult<T> & { rowCount: number }> {
    this.ensureNotDisposed();
    await this.ensureAuthContext();

    if (!this.authContext) {
      throw new GraphQLError(
        'Auth context not available. TenantAwareDBClient requires active authentication.',
        { extensions: { code: 'UNAUTHENTICATED' } },
      );
    }

    if (this.storage.getStore() && this.activeClient) {
      const result = await this.activeClient.query<T>(text, params);
      return { ...result, rowCount: result.rowCount ?? 0 };
    }

    return this.transaction(async client => {
      const result = await client.query<T>(text, params);
      return { ...result, rowCount: result.rowCount ?? 0 };
    });
  }

  /**
   * Execute a function within a transaction block.
   * Handles nested transactions using SAVEPOINTs.
   */
  public async transaction<T>(fn: (client: PoolClient) => Promise<T>): Promise<T> {
    this.ensureNotDisposed();
    await this.ensureAuthContext();

    if (!this.authContext) {
      throw new GraphQLError(
        'Auth context not available. TenantAwareDBClient requires active authentication.',
        { extensions: { code: 'UNAUTHENTICATED' } },
      );
    }

    if (this.storage.getStore()) {
      return this.executeTransactionInternal(fn);
    }

    return this.mutex.runExclusive(() => {
      this.ensureNotDisposed();
      return this.storage.run(true, () => {
        return this.executeTransactionInternal(fn);
      });
    });
  }

  private async executeTransactionInternal<T>(fn: (client: PoolClient) => Promise<T>): Promise<T> {
    // 1. Wait for initialization if in progress
    if (this.initializationPromise) {
      try {
        await this.initializationPromise;
      } catch {
        // Initialization failed.
        // We proceed to check (!this.activeClient) which will re-attempt or fail.
      }
    }

    // 2. Initialize if needed
    if (!this.activeClient) {
      this.initializationPromise = (async () => {
        const client = await this.dbProvider.pool.connect();
        try {
          await client.query('BEGIN');
          await this.setRLSVariables(client);
          this.activeClient = client;
        } catch (error) {
          client.release();
          throw error;
        }
      })();

      try {
        await this.initializationPromise;
      } finally {
        this.initializationPromise = null;
      }
    }

    // Guard: activeClient must be set by now
    if (!this.activeClient) {
      throw new Error('Failed to initialize database client');
    }

    const client = this.activeClient;
    this.transactionDepth++;
    let success = false;

    try {
      let result: T;

      // Use a savepoint for all nested scopes (depth > 1) to isolate failures
      // and allow partial success/failure within the shared transaction.
      if (this.transactionDepth > 1) {
        const savepointName = `sp_${this.transactionDepth}`;
        try {
          await client.query(`SAVEPOINT ${savepointName}`);
          result = await fn(client);
          await client.query(`RELEASE SAVEPOINT ${savepointName}`);
        } catch (error) {
          await client.query(`ROLLBACK TO SAVEPOINT ${savepointName}`);
          throw error;
        }
      } else {
        // Root scope (depth === 1) runs directly in the main transaction
        result = await fn(client);
      }

      success = true;
      return result;
    } catch (error) {
      // If we are the last active scope and an error occurred, we deliberately ROLLBACK.
      // Note: If depth > 1, the inner try/catch already handled the savepoint rollback,
      // so this block only handles the root transaction failure or unhandled critical errors.
      if (this.transactionDepth === 1) {
        // Ensure we don't try to rollback if already disposed or closed
        try {
          await client.query('ROLLBACK');
        } catch {
          // Ignore rollback errors (e.g. if connection closed)
        }
      }
      throw error;
    } finally {
      this.transactionDepth--;

      // If we are the last scope to exit, we are responsible for cleanup.
      if (this.transactionDepth === 0) {
        if (success && !this.isDisposed) {
          try {
            await client.query('COMMIT');
          } catch (commitError) {
            console.error('Failed to commit transaction:', commitError);
          }
        }
        client.release();
        this.activeClient = null;
      }
    }
  }

  /**
   * Set PostgreSQL session variables for Row-Level Security.
   */
  private async setRLSVariables(client: PoolClient): Promise<void> {
    if (!this.authContext) {
      throw new GraphQLError('Unauthenticated', {
        extensions: {
          code: 'UNAUTHENTICATED',
        },
      });
    }

    const { tenant, user, authType, activeReadScope } = this.authContext ?? {};

    // Write-target: the single business this request owns / writes to, derived
    // from the primary tenant business and the active scope. The auth context
    // already re-points `tenant.businessId` to this value; resolving it here
    // again keeps the RLS session correct as defense-in-depth.
    const businessIdValue = resolveWriteTargetBusinessId(tenant?.businessId, activeReadScope);

    if (!businessIdValue) {
      throw new Error('Missing businessId in AuthContext');
    }

    // API keys use a non-UUID identifier (e.g. "api-key:<id>") for app-level tracing.
    // The DB helper get_current_user_id() casts app.current_user_id to UUID and handles
    // empty string via NULLIF(..., ''), so we pass '' for API key sessions to avoid a
    // runtime cast error while explicitly clearing the setting.
    const userIdValue = authType === 'apiKey' ? '' : (user?.userId ?? null);

    // Read scope: the businesses this request may read from, serialized as a
    // Postgres array literal ('{uuid1,uuid2}') for get_current_business_scope().
    // When empty/absent we pass '' so the DB helper falls back to the single
    // write-target business. Writes remain pinned to app.current_business_id.
    const readScopeValue =
      activeReadScope && activeReadScope.businessIds.length > 0
        ? `{${activeReadScope.businessIds.map(id => `"${id.replace(/"/g, '\\"')}"`).join(',')}}`
        : '';

    await client.query(
      `
      SELECT
        set_config('app.current_business_id', $1, true),
        set_config('app.current_user_id', $2, true),
        set_config('app.auth_type', $3, true),
        set_config('app.current_business_scope', $4, true);
      `,
      [businessIdValue, userIdValue, authType, readScopeValue],
    );
  }

  /**
   * Manually dispose the client.
   */
  public async dispose(): Promise<void> {
    if (this.isDisposed) return;

    // In normal operation, activeClient should already be null because
    // executeTransactionInternal's finally block releases it.
    // If it's null, we can skip the expensive mutex acquisition.
    if (!this.activeClient) {
      this.isDisposed = true;
      return;
    }

    // If we get here, activeClient is not null, which means the transaction
    // didn't clean up properly. This is an abnormal situation (e.g., connection leak,
    // error in finally block, or dispose called prematurely).
    console.warn(
      'TenantAwareDBClient.dispose() called with activeClient still set. Forcing cleanup.',
    );

    // Use a timeout or race to prevent hanging indefinitely during disposal
    // If a query is stuck, we don't want to block the entire server request handler.
    const TIMEOUT_MS = 5000;
    let release: (() => void) | undefined;

    try {
      release = await Promise.race([
        this.mutex.acquire(),
        new Promise<() => void>((_, reject) =>
          setTimeout(() => reject(new Error('Timeout acquiring mutex')), TIMEOUT_MS),
        ),
      ]);
    } catch (e) {
      console.warn(
        'Timeout acquiring mutex during TenantAwareDBClient disposal. Connection may be in use.',
        e,
      );
      // Don't force cleanup - let executeTransactionInternal's finally block handle it.
      // Forcing cleanup here creates a race condition where we destroy a connection
      // that's actively being used, causing query failures.
      // However, mark as disposed to prevent further usage of this instance.
      this.isDisposed = true;
      return;
    }

    try {
      if (this.isDisposed) return;

      if (this.activeClient) {
        let hadRollbackError = false;
        try {
          // Attempt rollback, but catch errors if connection is busy/closed
          await Promise.race([
            this.activeClient.query('ROLLBACK'),
            new Promise((_, reject) =>
              setTimeout(() => reject(new Error('Rollback timeout')), 1000),
            ),
          ]);
        } catch (error) {
          hadRollbackError = true;
          console.error('Error disposing TenantAwareDBClient (rollback failed):', error);
        } finally {
          try {
            // Only destroy the connection (release(true)) if there was an error.
            // For normal disposal, return it to the pool (release()) to avoid pool exhaustion.
            // pg's release(true) removes the connection from the pool permanently.
            this.activeClient.release(hadRollbackError);
          } catch (e) {
            console.error('Error releasing client during disposal:', e);
          }
          this.activeClient = null;
        }
      }
      this.isDisposed = true;
    } finally {
      if (release) {
        release();
      }
    }
  }

  private ensureNotDisposed() {
    if (this.isDisposed) {
      throw new Error('TenantAwareDBClient is already disposed');
    }
  }

  /**
   * Lazy initialization of auth context on first use.
   * This ensures the async provider is called only when needed.
   */
  private async ensureAuthContext(): Promise<void> {
    if (this.authContextInitialized) {
      return;
    }
    if (!this.authContextProvider) {
      throw new GraphQLError(
        'Auth context not available. TenantAwareDBClient requires active authentication.',
        { extensions: { code: 'UNAUTHENTICATED' } },
      );
    }
    this.authContext = await this.authContextProvider.getAuthContext();
    this.authContextInitialized = true;
  }
}