@compas/store/src/query.js

Postgres & S3-compatible wrappers for common things

// @ts-nocheck

import {
  _compasSentryEnableQuerySpans,
  _compasSentryExport,
  isNil,
} from "@compas/stdlib";

/**
 * Format and append query parts, and execute the final result in a safe way.
 * Undefined values are skipped, as they are not allowed in queries.
 * The query call may be one of the interpolated values. Supports being called as a
 * template literal.
 *
 * @since 0.1.0
 *
 * @template T
 *
 * @param {TemplateStringsArray | Array<string>} strings
 * @param {...(import("../types/advanced-types.d.ts").QueryPartArg
 *   | Array<import("../types/advanced-types.d.ts").QueryPartArg>
 *   )} values
 * @returns {import("../types/advanced-types.d.ts").QueryPart<T>}
 */
export function query(strings, ...values) {
  /** @type {Array<string>} */
  let _strings = [];
  /** @type {Array<QueryPartArg>} */
  const _values = [];

  const result = /** @type {QueryPart<T>} */ {
    get strings() {
      return _strings;
    },
    get values() {
      return _values;
    },
    append,
    exec,
  };

  // Flatten nested query parts
  let didFlatten = false;
  for (let i = 0; i < strings.length - 1; ++i) {
    if (didFlatten) {
      didFlatten = false;
      _strings[_strings.length - 1] += strings[i];
    } else {
      _strings.push(strings[i]);
    }
    if (isQueryPart(values[i])) {
      append(values[i]);
      didFlatten = true;
    } else {
      _values.push(values[i]);
    }
  }

  if (didFlatten) {
    _strings[_strings.length - 1] += strings[strings.length - 1];
  } else {
    _strings.push(strings[strings.length - 1]);
  }

  return result;

  function append(query) {
    const last = _strings[_strings.length - 1];
    const [first, ...rest] = query.strings;
    _strings = [..._strings.slice(0, -1), `${last} ${first}`, ...rest];
    _values.push.apply(_values, query.values);
    return result;
  }

  /**
   * @param {import("../index.js").Postgres} sql
   * @returns {Promise<import("postgres").PendingQuery<T>>}
   */
  async function exec(sql) {
    let str = _strings[0];
    let valueIdx = 1;
    for (let i = 0; i < _values.length; ++i) {
      if (_values[i] === undefined) {
        str += `${_strings[i + 1]}`;
      } else {
        str += `$${valueIdx++}${_strings[i + 1]}`;
      }
    }

    // Strip out undefined values
    /** @type {NonNullable<QueryPartArg>} */
    const parameters = _values.filter((it) => it !== undefined);

    if (
      typeof _compasSentryExport?.startSpan === "function" &&
      _compasSentryEnableQuerySpans
    ) {
      return await _compasSentryExport.startSpan(
        {
          op: "db.query",
          name: str,
          attributes: {
            "db.system": "postgresql",
          },
          onlyIfParent: true,
        },
        () => sql.unsafe(str, parameters),
      );
    }

    return await sql.unsafe(str, parameters);
  }
}

/**
 * Check if the passed in value is an object generated by 'query``'.
 *
 * @since 0.1.0
 *
 * @param {any} query
 * @returns {query is import("../types/advanced-types.d.ts").QueryPart<any>}
 */
export function isQueryPart(query) {
  return (
    !isNil(query) &&
    Array.isArray(query?.strings) &&
    Array.isArray(query?.values)
  );
}

/**
 * Stringify a queryPart.
 * When interpolateParameters is true, we do a best effort in replacing the parameterized
 * query with the real params. If the result doesn't look right, please turn it off.
 *
 * @since 0.1.0
 *
 * @param {import("../types/advanced-types.d.ts").QueryPart<any>} queryPart
 * @param {{ interpolateParameters?: boolean }} options
 * @returns {string | {sql?: string, params?: Array<*>}}
 */
export function stringifyQueryPart(queryPart, { interpolateParameters } = {}) {
  if (!isQueryPart(queryPart)) {
    throw new Error(
      `'stringifyQueryPart' expects a query part produced by calling 'query\`\`'`,
    );
  }

  /** @type {string} */
  let sql = "";
  let params = [];
  queryPart.exec(
    /** @type {any} */ {
      unsafe(queryString, parameters) {
        sql = queryString.trim().replaceAll(/\s+/g, " ");
        params = /** @type {Array<any>} */ parameters;
      },
    },
  );

  if (!interpolateParameters) {
    return {
      sql,
      params,
    };
  }

  return sql.replace(/\$\d+/g, (match) => {
    const idx = parseInt(match.substring(1));
    const value = params[idx - 1];

    if (typeof value === "string") {
      return `'${value}'`;
    } else if (value instanceof Date) {
      return `'${value.toISOString()}'`;
    }

    return value;
  });
}

/**
 * Creates a transaction, executes the query, and rollback the transaction afterwards.
 * This is safe to use with insert, update and delete queries.
 *
 * By default returns text, but can also return json.
 * Note that explain output is highly depended on the current data and usage of the
 * tables.
 *
 * @since 0.1.0
 *
 * @param {import("../index.js").Postgres} sql
 * @param {import("../types/advanced-types.d.ts").QueryPart<any>} queryItem
 * @param {{ jsonResult?: boolean }} [options={}]
 * @returns {Promise<string|object>}
 */
export async function explainAnalyzeQuery(sql, queryItem, { jsonResult } = {}) {
  let result;

  try {
    await sql.begin(async (sql) => {
      if (jsonResult) {
        const intermediate =
          await query`EXPLAIN (ANALYZE, VERBOSE, BUFFERS, FORMAT JSON) ${queryItem}`.exec(
            sql,
          );
        result = intermediate[0];
      } else {
        const intermediate =
          await query`EXPLAIN (ANALYZE, VERBOSE, BUFFERS, FORMAT TEXT) ${queryItem}`.exec(
            sql,
          );

        result = intermediate.map((it) => it["QUERY PLAN"]).join("\n");
      }
      // Rollback the transaction
      throw new Error();
    });
  } catch {
    // Ignore
  }

  return result;
}