@bytebase/dbhub/dist/chunk-X5OCFRNL.js

Token-efficient database MCP server for PostgreSQL, MySQL, MariaDB, SQL Server, Oracle, SQLite

import {
  hasMutatingKeyword
} from "./chunk-55X2VY5D.js";
import {
  blankCommentsAndStrings
} from "./chunk-IXC4NQGC.js";

// src/utils/sql-row-limiter.ts
var SQLRowLimiter = class {
  /**
   * Check if a SQL statement is a row-returning query that can benefit from row limiting.
   *
   * Classification runs on the comment-blanked text, so a query introduced by a
   * `--` or `/* ... *\/` comment (a query tag, say) is classified by its real
   * leading keyword. The SQL sent to the server is never rewritten by this
   * check, so an attribution comment survives verbatim.
   *
   * `WITH` leads a CTE whose final statement is normally a SELECT, so it is
   * limitable too, except for a data-modifying CTE
   * (`WITH x AS (DELETE ... RETURNING *) SELECT ...`), where a LIMIT would cap
   * the rows handed back while the write still runs in full. Detection there is
   * the same keyword heuristic the read-only classifier uses, so a false
   * positive only means the statement keeps the old behaviour of not being limited.
   *
   * `dialect` selects the connector's scanner so dialect-specific quoting
   * (PostgreSQL dollar quotes, MySQL backticks, SQL Server brackets) is
   * blanked the same way the read-only classifier blanks it; without it only
   * ANSI syntax is recognized.
   */
  static isSelectQuery(sql, dialect) {
    const blankedSQL = blankCommentsAndStrings(sql, dialect).trim().toLowerCase();
    const firstKeyword = /^[(;\s]*([a-z_]+)/.exec(blankedSQL)?.[1] ?? "";
    if (firstKeyword === "select") {
      return true;
    }
    return firstKeyword === "with" && !hasMutatingKeyword(blankedSQL, dialect);
  }
  /**
   * Check if a SQL statement has a LIMIT clause of its own (not one nested in
   * a CTE body or subquery).
   */
  static hasLimitClause(sql, dialect) {
    return this.findTopLevelLimit(sql, dialect) !== null;
  }
  /**
   * Check if a SQL statement has a TOP clause of its own (SQL Server).
   */
  static hasTopClause(sql) {
    return this.findTopLevelTop(sql) !== null;
  }
  /**
   * Extract the statement's own LIMIT value. Null when it has no LIMIT of its
   * own, or when that LIMIT is a parameter placeholder rather than a literal
   * (see hasParameterizedLimit).
   */
  static extractLimitValue(sql, dialect) {
    return this.findTopLevelLimit(sql, dialect)?.value ?? null;
  }
  /**
   * Extract the statement's own TOP value (SQL Server), or null when it has none.
   */
  static extractTopValue(sql) {
    return this.findTopLevelTop(sql)?.value ?? null;
  }
  /**
   * Check if the statement's own LIMIT clause uses a parameter placeholder
   * ($1, ?, @p1) instead of a literal number.
   */
  static hasParameterizedLimit(sql, dialect) {
    const limit = this.findTopLevelLimit(sql, dialect);
    return limit !== null && limit.value === null;
  }
  /**
   * Add or tighten the LIMIT clause of a SQL statement
   */
  static applyLimitToQuery(sql, maxRows, dialect) {
    const limit = this.findTopLevelLimit(sql, dialect);
    if (limit !== null && limit.value !== null) {
      const effectiveLimit = Math.min(limit.value, maxRows);
      return `${sql.slice(0, limit.index)}LIMIT ${effectiveLimit}${sql.slice(limit.index + limit.length)}`;
    }
    const { sql: sqlWithoutSemicolon, semicolon } = trimSemicolon(sql);
    return `${sqlWithoutSemicolon}
LIMIT ${maxRows}${semicolon}`;
  }
  /**
   * Scan blanked (comment/string-free, length-preserving) SQL for parenthesis
   * depth, invoking onMatch for every regex hit at depth 0 (i.e. not nested
   * inside a subquery, function call, or window OVER clause).
   */
  static scanTopLevel(blankedSQL, regex, onMatch) {
    let depth = 0;
    let m;
    while ((m = regex.exec(blankedSQL)) !== null) {
      const token = m[0];
      if (token === "(") {
        depth++;
      } else if (token === ")") {
        depth--;
      } else if (depth === 0) {
        onMatch(m);
      }
    }
  }
  /**
   * The first or last regex match at parenthesis depth 0, i.e. one belonging
   * to the statement itself rather than to a nested query. Match indices refer
   * to positions in the original `sql`, since blankCommentsAndStrings
   * preserves length. The regex must offer `\(` and `\)` alternatives so depth
   * can be tracked, and must carry the `g` flag.
   */
  static findTopLevelMatch(sql, regex, pick, dialect) {
    let found = null;
    this.scanTopLevel(blankCommentsAndStrings(sql, dialect), regex, (m) => {
      if (pick === "last" || found === null) {
        found = m;
      }
    });
    return found;
  }
  /** The statement's own LIMIT clause, literal or parameterized ($1, ?, @p1). */
  static findTopLevelLimit(sql, dialect) {
    const match = this.findTopLevelMatch(
      sql,
      /\(|\)|\blimit\s+(?:(\d+)|\$\d+|\?|@p\d+)/gi,
      "last",
      dialect
    );
    if (match === null) {
      return null;
    }
    return {
      index: match.index,
      length: match[0].length,
      value: match[1] !== void 0 ? parseInt(match[1], 10) : null
    };
  }
  /**
   * The statement's own TOP clause (SQL Server), located on its own SELECT
   * (after an optional DISTINCT/ALL). `index`/`length` span the `TOP ...`
   * text only, so it can be spliced without disturbing `SELECT DISTINCT`.
   *
   * `value` is the clause's literal row count (`TOP 3`, `TOP (3)`), or null
   * when that number is not an upper bound on the rows returned: a
   * parenthesised expression or parameter (`TOP (@p1)`), `PERCENT`, or
   * `WITH TIES`. Such a TOP cannot simply be tightened; the statement has to
   * be wrapped and capped from the outside.
   */
  static findTopLevelTop(sql) {
    const match = this.findTopLevelMatch(
      sql,
      /\(|\)|\bselect\s+(?:(?:distinct|all)\s+)?(top)\b/gi,
      "first",
      "sqlserver"
    );
    if (match === null) {
      return null;
    }
    const blanked = blankCommentsAndStrings(sql, "sqlserver");
    const index = match.index + match[0].length - match[1].length;
    let end = match.index + match[0].length;
    let literal;
    const operand = /\s*(?:(\()|(\d+)\b)/y;
    operand.lastIndex = end;
    const operandMatch = operand.exec(blanked);
    if (operandMatch === null) {
      return null;
    }
    end = operand.lastIndex;
    if (operandMatch[1] !== void 0) {
      const close = this.findClosingParen(blanked, end - 1);
      if (close === -1) {
        return null;
      }
      const inner = blanked.slice(end, close);
      literal = /^\s*(\d+)\s*$/.exec(inner)?.[1];
      end = close + 1;
    } else {
      literal = operandMatch[2];
    }
    const modifiers = operandMatch[1] !== void 0 ? /(?:\s*percent\b)?(?:\s*with\s+ties\b)?/iy : /(?:\s+percent\b)?(?:\s+with\s+ties\b)?/iy;
    modifiers.lastIndex = end;
    const modifierMatch = modifiers.exec(blanked);
    const unbounded = modifierMatch !== null && modifierMatch[0].length > 0;
    end = modifiers.lastIndex;
    return {
      index,
      length: end - index,
      value: literal !== void 0 && !unbounded ? parseInt(literal, 10) : null
    };
  }
  /** Index of the `)` matching the `(` at `openIndex` in blanked SQL, or -1. */
  static findClosingParen(blankedSQL, openIndex) {
    let depth = 0;
    for (let i = openIndex; i < blankedSQL.length; i++) {
      const ch = blankedSQL[i];
      if (ch === "(") {
        depth++;
      } else if (ch === ")") {
        depth--;
        if (depth === 0) {
          return i;
        }
      }
    }
    return -1;
  }
  /**
   * The statement's own SELECT keyword (SQL Server), together with a
   * DISTINCT/ALL that follows it: T-SQL places TOP after those
   * (`SELECT DISTINCT TOP n`), so that is where a TOP gets inserted. For a
   * CTE this is the final SELECT, not the one inside a CTE body.
   */
  static findTopLevelSelect(sql) {
    return this.findTopLevelMatch(sql, /\(|\)|\bselect\b(?:\s+(?:distinct|all)\b)?/gi, "first", "sqlserver");
  }
  /**
   * The statement's own `OFFSET ... ROWS [FETCH NEXT n ROWS ONLY]` clause
   * (SQL Server). T-SQL forbids TOP in a query that has OFFSET, so such a
   * statement is capped through its FETCH count instead.
   *
   * `fetch` is the FETCH count (null when the clause has no FETCH);
   * `fetchInsertIndex` is where a FETCH can be appended, or -1 when the
   * OFFSET clause has a shape this scanner does not parse (an arithmetic
   * expression, say), in which case the statement is wrapped.
   */
  static findTopLevelOffset(sql) {
    const offset = this.findTopLevelMatch(sql, /\(|\)|\boffset\b/gi, "first", "sqlserver");
    if (offset === null) {
      return null;
    }
    const blanked = blankCommentsAndStrings(sql, "sqlserver");
    const clause = /(offset\s+(?:\d+|@\w+|\?)\s+rows?\b)(?:(\s+fetch\s+(?:first|next)\s+)((\d+)|@\w+|\?)\s+rows?\s+only\b)?/iy;
    clause.lastIndex = offset.index;
    const match = clause.exec(blanked);
    if (match === null) {
      return { fetch: null, fetchInsertIndex: -1 };
    }
    if (match[2] === void 0) {
      const end = match.index + match[0].length;
      const unparsedFetch = /\s*fetch\b/iy;
      unparsedFetch.lastIndex = end;
      if (unparsedFetch.test(blanked)) {
        return { fetch: null, fetchInsertIndex: -1 };
      }
      return { fetch: null, fetchInsertIndex: end };
    }
    return {
      fetch: {
        index: match.index + match[1].length + match[2].length,
        length: match[3].length,
        value: match[4] !== void 0 ? parseInt(match[4], 10) : null
      },
      fetchInsertIndex: -1
    };
  }
  /**
   * Check if a SQL statement combines multiple SELECTs with a set operator
   * (UNION [ALL], INTERSECT, EXCEPT) at the top level — i.e. not nested
   * inside a subquery already wrapped in parentheses.
   */
  static hasSetOperator(sql) {
    return this.findTopLevelMatch(sql, /\(|\)|\bunion\b|\bintersect\b|\bexcept\b/gi, "first", "sqlserver") !== null;
  }
  /**
   * Find the start index of a top-level trailing ORDER BY clause (not one
   * nested inside a subquery or a window function's OVER (...) clause).
   * Returns -1 if none exists.
   */
  static findTopLevelOrderByIndex(sql) {
    return this.findTopLevelMatch(sql, /\(|\)|\border\s+by\b/gi, "last", "sqlserver")?.index ?? -1;
  }
  /**
   * Add or modify TOP clause in a SQL statement (SQL Server)
   */
  static applyTopToQuery(sql, maxRows) {
    const offset = this.findTopLevelOffset(sql);
    if (offset !== null) {
      if (offset.fetch !== null && offset.fetch.value !== null) {
        const effectiveFetch = Math.min(offset.fetch.value, maxRows);
        return `${sql.slice(0, offset.fetch.index)}${effectiveFetch}${sql.slice(offset.fetch.index + offset.fetch.length)}`;
      }
      if (offset.fetch === null && offset.fetchInsertIndex !== -1) {
        return `${sql.slice(0, offset.fetchInsertIndex)} FETCH NEXT ${maxRows} ROWS ONLY${sql.slice(offset.fetchInsertIndex)}`;
      }
      return this.wrapWithTop(sql, maxRows, false);
    }
    if (this.hasSetOperator(sql)) {
      return this.wrapWithTop(sql, maxRows, true);
    }
    const existingTop = this.findTopLevelTop(sql);
    if (existingTop !== null) {
      if (existingTop.value === null) {
        return this.wrapWithTop(sql, maxRows, false);
      }
      const effectiveTop = Math.min(existingTop.value, maxRows);
      return `${sql.slice(0, existingTop.index)}TOP ${effectiveTop}${sql.slice(existingTop.index + existingTop.length)}`;
    }
    const selectMatch = this.findTopLevelSelect(sql);
    if (selectMatch === null) {
      return sql;
    }
    const insertAt = selectMatch.index + selectMatch[0].length;
    return `${sql.slice(0, insertAt)} TOP ${maxRows}${sql.slice(insertAt)}`;
  }
  /**
   * Cap a statement from the outside: `SELECT TOP n * FROM (<body>) AS subq`.
   *
   * A leading CTE stays outside the derived table: T-SQL has no
   * `SELECT ... FROM (WITH ...) AS subq` form, but a CTE declared before the
   * SELECT is still in scope inside that derived table.
   *
   * T-SQL disallows ORDER BY inside a derived table unless that derived table
   * itself has TOP/OFFSET/FOR XML. With `hoistOrderBy` a top-level trailing
   * ORDER BY moves outside the wrap (the set-operator case, where no TOP
   * belongs to the body as a whole); without it the ORDER BY stays inside,
   * which callers use when the body's own TOP/OFFSET both legalises it and
   * depends on it.
   */
  static wrapWithTop(sql, maxRows, hoistOrderBy) {
    const { sql: sqlWithoutSemicolon, semicolon } = trimSemicolon(sql);
    const cteIndex = this.findTopLevelSelect(sqlWithoutSemicolon)?.index ?? 0;
    const ctePrefix = sqlWithoutSemicolon.slice(0, cteIndex);
    let body = sqlWithoutSemicolon.slice(cteIndex);
    const optionIndex = this.findTopLevelOptionIndex(body);
    const optionClause = optionIndex === -1 ? "" : ` ${body.slice(optionIndex).trim()}`;
    if (optionIndex !== -1) {
      body = body.slice(0, optionIndex).trimEnd();
    }
    const orderByIndex = hoistOrderBy ? this.findTopLevelOrderByIndex(body) : -1;
    if (orderByIndex !== -1) {
      const innerSql = body.slice(0, orderByIndex).trimEnd();
      const orderByClause = body.slice(orderByIndex).trim();
      return `${ctePrefix}SELECT TOP ${maxRows} * FROM (${innerSql}
) AS subq ${orderByClause}${optionClause}${semicolon}`;
    }
    return `${ctePrefix}SELECT TOP ${maxRows} * FROM (${body}
) AS subq${optionClause}${semicolon}`;
  }
  /** Start index of the statement's own trailing `OPTION (...)` query hint, or -1. */
  static findTopLevelOptionIndex(sql) {
    return this.findTopLevelMatch(sql, /\(|\)|\boption\s*(?=\()/gi, "last", "sqlserver")?.index ?? -1;
  }
  /**
   * Apply maxRows limit to a row-returning query only
   *
   * This method is used by PostgreSQL, MySQL, MariaDB, and SQLite connectors which all support
   * the LIMIT clause syntax. SQL Server uses applyMaxRowsForSQLServer() instead with TOP syntax.
   *
   * For parameterized LIMIT clauses (e.g., LIMIT $1 or LIMIT ?), we wrap the query in a subquery
   * to enforce max_rows as a hard cap, since the parameter value is not known until runtime.
   *
   * `dialect` selects the connector's comment/string scanner (see isSelectQuery).
   */
  static applyMaxRows(sql, maxRows, dialect) {
    if (!maxRows || !this.isSelectQuery(sql, dialect)) {
      return sql;
    }
    if (this.hasParameterizedLimit(sql, dialect)) {
      const { sql: sqlWithoutSemicolon, semicolon } = trimSemicolon(sql);
      return `SELECT * FROM (${sqlWithoutSemicolon}
) AS subq LIMIT ${maxRows}${semicolon}`;
    }
    return this.applyLimitToQuery(sql, maxRows, dialect);
  }
  /**
   * Apply maxRows limit to a row-returning query using SQL Server TOP syntax
   */
  static applyMaxRowsForSQLServer(sql, maxRows) {
    if (!maxRows || !this.isSelectQuery(sql, "sqlserver")) {
      return sql;
    }
    return this.applyTopToQuery(sql, maxRows);
  }
  /**
   * Apply maxRows to a row-returning query using Oracle's row-limiting clause.
   *
   * The statement is always wrapped in an inline view and capped from the
   * outside: `SELECT * FROM (<sql>) FETCH FIRST n ROWS ONLY`. Wrapping rather
   * than splicing keeps every shape correct with one rule - a set operation
   * (UNION & co.) is capped as a whole, a trailing ORDER BY stays inside the
   * view where Oracle honours it, and a query that already has its own
   * `FETCH FIRST` / `ROWNUM` cap simply returns the smaller of the two. Oracle
   * also accepts a leading CTE (subquery factoring) inside an inline view, so
   * `WITH ... SELECT` needs no special casing.
   */
  static applyMaxRowsForOracle(sql, maxRows) {
    if (!maxRows || !this.isSelectQuery(sql, "oracle") || this.hasOracleForUpdate(sql)) {
      return sql;
    }
    const { sql: sqlWithoutSemicolon } = trimSemicolon(sql);
    return `SELECT * FROM (${sqlWithoutSemicolon}
) FETCH FIRST ${maxRows} ROWS ONLY`;
  }
  /**
   * A statement-level `FOR UPDATE` (locking read). Oracle allows neither
   * `FOR UPDATE` inside an inline view nor a row-limiting clause alongside
   * it, so such a statement cannot be capped and is left as written.
   */
  static hasOracleForUpdate(sql) {
    return this.findTopLevelMatch(sql, /\(|\)|\bfor\s+update\b/gi, "first", "oracle") !== null;
  }
  /**
   * Oracle variant of applyMaxRowsWithTruncationProbe. Because the cap always
   * wraps (see applyMaxRowsForOracle), the probe is applied to every
   * row-returning statement; a query whose own cap is tighter than maxRows
   * returns fewer than maxRows + 1 rows and is therefore never flagged.
   * A `FOR UPDATE` statement is never capped (see hasOracleForUpdate).
   */
  static applyMaxRowsForOracleWithTruncationProbe(sql, maxRows) {
    if (!maxRows || !this.isSelectQuery(sql, "oracle") || this.hasOracleForUpdate(sql)) {
      return { sql, probeApplied: false };
    }
    return { sql: this.applyMaxRowsForOracle(sql, maxRows + 1), probeApplied: true };
  }
  /**
   * Like applyMaxRows, but requests one probe row (maxRows + 1) whenever the
   * cap is the binding constraint, so callers can detect truncation exactly
   * instead of guessing from `rowCount === maxRows` (which a table with
   * exactly maxRows rows also produces).
   *
   * `probeApplied` is false when the cap cannot fire: no maxRows, not a
   * SELECT, or the query's own literal LIMIT is already within the cap (the
   * user asked for fewer rows than the cap — that is their limit, not
   * truncation).
   */
  static applyMaxRowsWithTruncationProbe(sql, maxRows, dialect) {
    if (!maxRows || !this.isSelectQuery(sql, dialect)) {
      return { sql, probeApplied: false };
    }
    if (!this.hasParameterizedLimit(sql, dialect)) {
      const existingLimit = this.extractLimitValue(sql, dialect);
      if (existingLimit !== null && existingLimit <= maxRows) {
        return { sql, probeApplied: false };
      }
    }
    return { sql: this.applyMaxRows(sql, maxRows + 1, dialect), probeApplied: true };
  }
  /**
   * SQL Server variant of applyMaxRowsWithTruncationProbe, using TOP syntax.
   * A set-operator statement is always wrapped (any TOP on an individual
   * branch caps only that branch, not the combined output), so the probe
   * applies there regardless of branch-level TOPs.
   */
  static applyMaxRowsForSQLServerWithTruncationProbe(sql, maxRows) {
    if (!maxRows || !this.isSelectQuery(sql, "sqlserver")) {
      return { sql, probeApplied: false };
    }
    const ownCap = this.findTopLevelOffset(sql)?.fetch?.value ?? (this.hasSetOperator(sql) ? null : this.extractTopValue(sql));
    if (ownCap !== null && ownCap <= maxRows) {
      return { sql, probeApplied: false };
    }
    return { sql: this.applyMaxRowsForSQLServer(sql, maxRows + 1), probeApplied: true };
  }
  /**
   * Post-process a result set produced by a truncation-probe rewrite: when
   * the probe row came back (more than maxRows rows), drop it, clamp the
   * count to maxRows, and mark the set truncated so consumers can tell a
   * capped result from a complete one.
   */
  static flagTruncation(resultSet, maxRows, probeApplied) {
    if (!probeApplied || !maxRows || resultSet.rows.length <= maxRows) {
      return;
    }
    resultSet.rows = resultSet.rows.slice(0, maxRows);
    resultSet.rowCount = maxRows;
    resultSet.truncated = true;
  }
};
function trimSemicolon(sql) {
  const trimmed = sql.trim();
  return trimmed.endsWith(";") ? { sql: trimmed.slice(0, -1), semicolon: ";" } : { sql: trimmed, semicolon: "" };
}

// src/utils/resource-cleanup.ts
async function closeQuietly(close) {
  try {
    await close();
  } catch {
  }
}

export {
  SQLRowLimiter,
  closeQuietly
};