market_screener/server/db/QueryBuilder.ts

/**
 * Type-safe query builder for SQLite.
 *
 * Prevents SQL injection by:
 * 1. Enforcing parameterized queries (? placeholders)
 * 2. Building SQL dynamically only for schema-safe values (table/column names are validated against a whitelist)
 * 3. Keeping all user input in parameter arrays, never in the SQL string
 *
 * Usage:
 *   const qb = new QueryBuilder('holdings');
 *   qb.select(['ticker', 'shares']).where('type = ?', ['stock']).orderBy('ticker');
 *   const stmt = db.prepare(qb.build());
 *   stmt.all(...qb.params());
 */

type QueryType = 'SELECT' | 'INSERT' | 'UPDATE' | 'DELETE';

interface WhereClause {
  expression: string;
  params: unknown[];
}

/**
 * Whitelist of safe column and table names.
 * Prevents injection via column/table names.
 */
const SAFE_COLUMNS = new Set([
  // holdings table
  'ticker',
  'shares',
  'cost_basis',
  'type',
  'source',
  // market_calls table
  'id',
  'title',
  'quarter',
  'date',
  'thesis',
  'tickers',
  'snapshot',
  'created_at',
]);

const SAFE_TABLES = new Set(['holdings', 'market_calls']);

/**
 * Validates a column name against the whitelist.
 * Throws if not in whitelist to prevent column name injection.
 */
function validateColumn(col: string): void {
  if (!SAFE_COLUMNS.has(col.toLowerCase())) {
    throw new Error(`Unsafe column name: ${col}. Only whitelisted columns allowed.`);
  }
}

/**
 * Validates a table name against the whitelist.
 * Throws if not in whitelist to prevent table name injection.
 */
function validateTable(table: string): void {
  if (!SAFE_TABLES.has(table.toLowerCase())) {
    throw new Error(`Unsafe table name: ${table}. Only whitelisted tables allowed.`);
  }
}

/**
 * QueryBuilder — type-safe, injectable-resistant query construction.
 */
export class QueryBuilder {
  private type: QueryType | null = null;
  private table: string;
  private selectCols: string[] = [];
  private whereClausesList: WhereClause[] = [];
  private orderByCols: { col: string; direction: 'ASC' | 'DESC' }[] = [];
  private limitVal: number | null = null;
  private offsetVal: number | null = null;

  // For INSERT
  private insertCols: string[] = [];
  private insertParamCount = 0;

  // For UPDATE
  private updateAssignments: { col: string; paramIndex: number }[] = [];

  private allParams: unknown[] = [];

  constructor(table: string) {
    validateTable(table);
    this.table = table;
  }

  /**
   * SELECT query builder.
   * Columns are validated against whitelist.
   */
  select(columns: string[]): this {
    if (this.type !== null) throw new Error('Query type already set');
    this.type = 'SELECT';
    for (const col of columns) {
      validateColumn(col);
      this.selectCols.push(col);
    }
    return this;
  }

  /**
   * INSERT query builder.
   * Columns are validated; values go into parameter array.
   */
  insert(columns: string[], values: unknown[]): this {
    if (this.type !== null) throw new Error('Query type already set');
    if (columns.length !== values.length) {
      throw new Error('Column/value count mismatch');
    }
    this.type = 'INSERT';
    for (const col of columns) {
      validateColumn(col);
      this.insertCols.push(col);
    }
    this.insertParamCount = values.length;
    this.allParams.push(...values);
    return this;
  }

  /**
   * UPDATE query builder.
   * Column names validated; values go into parameter array.
   */
  update(updates: Record<string, unknown>): this {
    if (this.type !== null) throw new Error('Query type already set');
    this.type = 'UPDATE';
    let paramIndex = 0;
    for (const [col, value] of Object.entries(updates)) {
      validateColumn(col);
      this.updateAssignments.push({ col, paramIndex });
      this.allParams.push(value);
      paramIndex++;
    }
    return this;
  }

  /**
   * DELETE query builder.
   */
  delete(): this {
    if (this.type !== null) throw new Error('Query type already set');
    this.type = 'DELETE';
    return this;
  }

  /**
   * WHERE clause(s).
   * Expression is NOT validated (it should be safe from app logic);
   * params are added to the parameter array.
   *
   * Example: .where('type = ? AND shares > ?', ['stock', 10])
   */
  where(expression: string, params: unknown[] = []): this {
    this.whereClausesList.push({ expression, params });
    this.allParams.push(...params);
    return this;
  }

  /**
   * ORDER BY clause.
   * Column names are validated.
   */
  orderBy(column: string, direction: 'ASC' | 'DESC' = 'ASC'): this {
    validateColumn(column);
    this.orderByCols.push({ col: column, direction });
    return this;
  }

  /**
   * LIMIT clause.
   */
  limit(count: number): this {
    if (count < 0) throw new Error('LIMIT must be non-negative');
    this.limitVal = count;
    return this;
  }

  /**
   * OFFSET clause.
   */
  offset(count: number): this {
    if (count < 0) throw new Error('OFFSET must be non-negative');
    this.offsetVal = count;
    return this;
  }

  /**
   * Build the final SQL string.
   * The query is built dynamically but with no injection points:
   * - Table/column names from whitelist only
   * - All user input in the parameter array
   */
  build(): string {
    if (this.type === null) throw new Error('Query type not set');

    let sql = '';

    switch (this.type) {
      case 'SELECT': {
        const cols = this.selectCols.length > 0 ? this.selectCols.join(', ') : '*';
        sql = `SELECT ${cols} FROM ${this.table}`;
        break;
      }

      case 'INSERT': {
        const cols = this.insertCols.join(', ');
        const placeholders = Array(this.insertParamCount).fill('?').join(', ');
        sql = `INSERT INTO ${this.table} (${cols}) VALUES (${placeholders})`;
        break;
      }

      case 'UPDATE': {
        const assignments = this.updateAssignments.map((a) => `${a.col} = ?`).join(', ');
        sql = `UPDATE ${this.table} SET ${assignments}`;
        break;
      }

      case 'DELETE': {
        sql = `DELETE FROM ${this.table}`;
        break;
      }
    }

    // Add WHERE clause(s)
    if (this.whereClausesList.length > 0) {
      const whereExpressions = this.whereClausesList.map((w) => `(${w.expression})`).join(' AND ');
      sql += ` WHERE ${whereExpressions}`;
    }

    // Add ORDER BY
    if (this.orderByCols.length > 0) {
      const orderExpressions = this.orderByCols.map((o) => `${o.col} ${o.direction}`).join(', ');
      sql += ` ORDER BY ${orderExpressions}`;
    }

    // Add LIMIT
    if (this.limitVal !== null) {
      sql += ` LIMIT ${this.limitVal}`;
    }

    // Add OFFSET
    if (this.offsetVal !== null) {
      sql += ` OFFSET ${this.offsetVal}`;
    }

    return sql;
  }

  /**
   * Return the accumulated parameter array.
   * This is what gets passed to db.prepare(...).run(...params).
   */
  params(): unknown[] {
    return this.allParams;
  }
}