@uwdata/mosaic-sql 0.16.2 → 0.18.0

package/src/ast/{query.js → query.ts}

@@ -1,10 +1,8 @@
-/**
- * @import { FilterExpr, FromExpr, GroupByExpr, OrderByExpr, SelectExpr, WithExpr } from '../types.js'
- * @import { SampleMethod } from './sample.js'
- */
+import type { FilterExpr, FromExpr, GroupByExpr, MaybeArray, OrderByExpr, SelectExpr, WithExpr } from '../types.js';
+import type { SampleMethod } from './sample.js';
 import { DESCRIBE_QUERY, SELECT_QUERY, SET_OPERATION } from '../constants.js';
-import { asNode, asTableRef, asVerbatim } from '../util/ast.js';
-import { exprList } from '../util/function.js';
+import { asNode, asVerbatim, maybeTableRef } from '../util/ast.js';
+import { exprList, nodeList } from '../util/function.js';
 import { unquote } from '../util/string.js';
 import { isArray, isString } from '../util/type-check.js';
 import { isColumnRef } from './column-ref.js';
@@ -13,153 +11,143 @@ import { ExprNode, SQLNode, isNode } from './node.js';
 import { SampleClauseNode } from './sample.js';
 import { SelectClauseNode } from './select.js';
 import { isTableRef } from './table-ref.js';
-import { WindowClauseNode } from './window.js';
+import { WindowClauseNode, type WindowDefNode } from './window.js';
 import { WithClauseNode } from './with.js';
 
 /**
  * Check if a value is a selection query or set operation.
- * @param {*} value The value to check.
- * @returns {value is Query}
+ * @param value The value to check.
  */
-export function isQuery(value) {
+export function isQuery(value: unknown): value is Query {
   return value instanceof Query;
 }
 
 /**
  * Check if a value is a selection query.
- * @param {*} value The value to check.
- * @returns {value is SelectQuery}
+ * @param value The value to check.
  */
-export function isSelectQuery(value) {
+export function isSelectQuery(value: unknown): value is SelectQuery {
   return value instanceof SelectQuery;
 }
 
 /**
  * Check if a value is a describe query.
- * @param {*} value The value to check.
- * @returns {value is DescribeQuery}
+ * @param value The value to check.
  */
-export function isDescribeQuery(value) {
+export function isDescribeQuery(value: unknown): value is DescribeQuery {
   return value instanceof DescribeQuery;
 }
 
 export class Query extends ExprNode {
   /**
    * Create a new WITH clause with the given CTE queries.
-   * @param {...WithExpr} expr The WITH CTE queries.
-   * @returns {WithClause}
+   * @param expr The WITH CTE queries.
    */
-  static with(...expr) {
+  static with(...expr: WithExpr[]) {
     return new WithClause(...expr);
   }
 
   /**
    * Create a new select query with the given SELECT expressions.
-   * @param {...SelectExpr} expr The SELECT expressions.
-   * @returns {SelectQuery}
+   * @param expr The SELECT expressions.
    */
-  static select(...expr) {
+  static select(...expr: SelectExpr[]) {
     return new SelectQuery().select(...expr);
   }
 
   /**
    * Create a new select query with the given FROM expressions.
-   * @param {...FromExpr} expr The FROM expressions.
-   * @returns {SelectQuery}
+   * @param expr The FROM expressions.
    */
-  static from(...expr) {
+  static from(...expr: FromExpr[]) {
     return new SelectQuery().from(...expr);
   }
 
   /**
    * Create a new UNION set operation over the given queries.
-   * @param {...Query} queries The queries.
-   * @returns {SetOperation}
+   * @param queries The queries.
    */
-  static union(...queries) {
+  static union(...queries: MaybeArray<Query>[]) {
     return new SetOperation('UNION', queries.flat());
   }
 
   /**
    * Create a new UNION ALL set operation over the given queries.
-   * @param {...Query} queries The queries.
-   * @returns {SetOperation}
+   * @param queries The queries.
    */
-  static unionAll(...queries) {
+  static unionAll(...queries: MaybeArray<Query>[]) {
     return new SetOperation('UNION ALL', queries.flat());
   }
 
   /**
    * Create a new INTERSECT set operation over the given queries.
-   * @param {...Query} queries The queries.
-   * @returns {SetOperation}
+   * @param queries The queries.
    */
-  static intersect(...queries) {
+  static intersect(...queries: MaybeArray<Query>[]) {
     return new SetOperation('INTERSECT', queries.flat());
   }
 
   /**
    * Create a new EXCEPT set operation over the given queries.
-   * @param {...Query} queries The queries.
-   * @returns {SetOperation}
+   * @param queries The queries.
    */
-  static except(...queries) {
+  static except(...queries: MaybeArray<Query>[]) {
     return new SetOperation('EXCEPT', queries.flat());
   }
 
   /**
    * Create a new describe query for the given input query.
-   * @param {Query} query The query to describe.
-   * @returns {DescribeQuery}
+   * @param query The query to describe.
    */
-  static describe(query) {
+  static describe(query: Query) {
     return new DescribeQuery(query);
   }
 
+  _with: WithClauseNode[] = [];
+  _orderby: ExprNode[] = [];
+  _limitPerc: boolean = false;
+  _limit?: ExprNode;
+  _offset?: ExprNode;
+  cteFor?: Query | null = null;
+
   /**
    * Instantiate a new query.
    */
-  constructor(type) {
+  constructor(type: string) {
     super(type);
-    /** @type {WithClauseNode[]} */
-    this._with = [];
-    /** @type {ExprNode[]} */
-    this._orderby = [];
-    /** @type {number} */
-    this._limit = undefined;
-    /** @type {number} */
-    this._offset = undefined;
-    /** @type {Query | null} */
-    this.cteFor = null;
   }
 
   /**
    * Return a list of subqueries.
-   * @returns {Query[]}
    */
-  get subqueries() {
+  get subqueries(): Query[] {
     return [];
   }
 
   /**
    * Clone this query.
-   * @returns {Query}
    */
-  clone() {
+  clone(): this {
     return this;
   }
 
+  /**
+   * Add a pointer to the query for which this query is a CTE.
+   * @param query The query for which this query is a CTE.
+   */
+  setCteFor(query: Query | null): void {
+    this.cteFor = query;
+  }
+
   /**
    * Add WITH common table expressions (CTEs).
-   * @param  {...WithExpr} expr Expressions to add.
-   * @returns {this}
+   * @param expr Expressions to add.
    */
-  with(...expr) {
-    /** @type {WithClauseNode[]} */
-    const list = [];
-    const add = (name, q) => {
+  with(...expr: WithExpr[]): this {
+    const list: WithClauseNode[] = [];
+    const add = (name: string, q: Query) => {
       const query = q.clone();
-      query.cteFor = this;
+      query.setCteFor(this);
       list.push(new WithClauseNode(name, query));
     };
     expr.flat().forEach(e => {
@@ -172,72 +160,76 @@ export class Query extends ExprNode {
 
   /**
    * Add ORDER BY expressions.
-   * @param  {...OrderByExpr} expr Expressions to add.
-   * @returns
+   * @param expr Expressions to add.
    */
-  orderby(...expr) {
-    this._orderby = this._orderby.concat(exprList(expr));
+  orderby(...expr: OrderByExpr[]): this {
+    this._orderby = this._orderby.concat(nodeList(expr));
+    return this;
+  }
+
+  /**
+   * Set the query result LIMIT as a percentage value.
+   * @param value The limit percentage value.
+   */
+  limitPercent(value: number | ExprNode): this {
+    this._limitPerc = true;
+    this._limit = asNode(value);
     return this;
   }
 
   /**
    * Set the query result LIMIT.
-   * @param {number} value The limit value.
-   * @returns {this}
+   * @param value The limit value.
    */
-  limit(value) {
-    this._limit = Number.isFinite(value) ? value : undefined;
+  limit(value: number | ExprNode): this {
+    this._limitPerc = false;
+    this._limit = asNode(value);
     return this;
   }
 
   /**
    * Set the query result OFFSET.
-   * @param {number} value The offset value.
-   * @returns {this}
+   * @param value The offset value.
    */
-  offset(value) {
-    this._offset = Number.isFinite(value) ? value : undefined;
+  offset(value: number | ExprNode): this {
+    this._offset = asNode(value);
     return this;
   }
 }
 
 export class SelectQuery extends Query {
+  _select: SelectClauseNode[] = [];
+  _from: FromClauseNode[] = [];
+  _where: ExprNode[] = [];
+  _sample?: SampleClauseNode;
+  _groupby: ExprNode[] = [];
+  _having: ExprNode[] = [];
+  _window: WindowClauseNode[] = [];
+  _qualify: ExprNode[] = [];
+  _distinct: boolean = false;
+
   /**
    * Instantiate a new select query.
    */
   constructor() {
     super(SELECT_QUERY);
-    /** @type {SelectClauseNode[]} */
-    this._select = [];
-    /** @type {FromClauseNode[]} */
-    this._from = [];
-    /** @type {ExprNode[]} */
-    this._where = [];
-    /** @type {SampleClauseNode} */
-    this._sample = undefined;
-    /** @type {ExprNode[]} */
-    this._groupby = [];
-    /** @type {ExprNode[]} */
-    this._having = [];
-    /** @type {WindowClauseNode[]} */
-    this._window = [];
-    /** @type {ExprNode[]} */
-    this._qualify = [];
   }
 
   /**
    * Return a list of subqueries.
-   * @returns {Query[]}
    */
-  get subqueries() {
+  get subqueries(): Query[] {
     // build map of ctes within base query WITH clause
     const q = this.cteFor || this;
     const w = q instanceof SelectQuery ? q._with : [];
-    const cte = w.reduce((obj, c) => (obj[c.name] = c.query, obj), {});
+    const cte = w.reduce(
+      (obj, c) => (obj[c.name] = c.query, obj),
+      {} as Record<string, Query>
+    );
 
     // extract subqueries in FROM clause
     // unused CTEs will be ignored
-    const queries = [];
+    const queries: Query[] = [];
     this._from.forEach(({ expr }) => {
       if (isQuery(expr)) {
         queries.push(expr);
@@ -251,68 +243,70 @@ export class SelectQuery extends Query {
 
   /**
    * Clone this query.
-   * @returns {SelectQuery}
    */
-  clone() {
+  clone(): this {
     return Object.assign(new SelectQuery(), this);
   }
 
   /**
    * Add SELECT expressions.
-   * @param {...SelectExpr} expr Expressions to add.
-   * @returns {this}
-   */
-  select(...expr) {
-    /** @type {SelectClauseNode[]} */
-    const list = [];
-    const add = (v, as) => list.push(
-      new SelectClauseNode(v == null ? v : asNode(v), unquote(as))
-    );
+   * @param expr Expressions to add.
+   */
+  select(...expr: SelectExpr[]): this {
+    const keys: Set<string> = new Set;
+    const list: SelectClauseNode[] = [];
+
+    const add = (v: unknown, as: string) => {
+      const key = unquote(as)!;
+      keys.add(key);
+      if (v) list.push(new SelectClauseNode(asNode(v), key));
+    };
+
     expr.flat().forEach(e => {
       if (e == null) return;
       else if (isString(e)) add(e, e);
       else if (isColumnRef(e)) add(e, e.column);
       else if (isArray(e)) add(e[1], e[0]);
+      else if (isNode(e)) add(e, `${e}`);
       else for (const alias in e) add(e[alias], alias);
     });
 
-    const keys = new Set(list.map(x => x.alias));
-    this._select = this._select
-      .filter(x => !keys.has(x.alias))
-      .concat(list.filter(x => x.expr));
+    this._select = this._select.filter(x => !keys.has(x.alias)).concat(list);
     return this;
   }
 
   /**
    * Set SELECT expressions, replacing any prior expressions.
-   * @param {...SelectExpr} expr Expressions to add.
-   * @returns {this}
+   * @param expr Expressions to add.
    */
-  setSelect(...expr) {
+  setSelect(...expr: SelectExpr[]): this {
     this._select = [];
     return this.select(...expr);
   }
 
   /**
    * Indicate if this query should retrieve distinct values only.
-   * @param {boolean} value The distinct flag
-   * @returns {this}
+   * @param value The distinct flag
    */
-  distinct(value = true) {
+  distinct(value: boolean = true): this {
     this._distinct = !!value;
     return this;
   }
 
   /**
    * Add table FROM expressions.
-   * @param  {...FromExpr} expr Expressions to add.
-   * @returns {this}
+   * @param expr Expressions to add.
    */
-  from(...expr) {
-    const list = [];
-    const add = (v, as) => list.push(new FromClauseNode(asTableRef(v), unquote(as)));
+  from(...expr: FromExpr[]): this {
+    const list: FromClauseNode[] = [];
+
+    const add = (v: string | string[] | SQLNode, as?: string) => {
+      list.push(new FromClauseNode(maybeTableRef(v), unquote(as)));
+    };
+
     expr.flat().forEach(e => {
       if (e == null) return;
+      else if (e instanceof FromClauseNode) list.push(e);
       else if (isString(e)) add(e, e);
       else if (isTableRef(e)) add(e, e.name);
       else if (isNode(e)) add(e);
@@ -325,24 +319,25 @@ export class SelectQuery extends Query {
 
   /**
    * Set FROM expressions, replacing any prior expressions.
-   * @param {...FromExpr} expr Expressions to add.
-   * @returns {this}
+   * @param expr Expressions to add.
    */
-  setFrom(...expr) {
+  setFrom(...expr: FromExpr[]): this {
     this._from = [];
     return this.from(...expr);
   }
 
   /**
    * Set SAMPLE settings.
-   * @param {number | SampleClauseNode} value Either a sample clause node
-   *  or the sample size as either a row count or percentage.
-   * @param {SampleMethod} [method] The sampling method
-   *  to use.
-   * @param {number} [seed] The random seed.
-   * @returns {this}
-   */
-  sample(value, method, seed) {
+   * @param value Either a sample clause node or the sample size as either
+   *  a row count or percentage.
+   * @param method The sampling method to use.
+   * @param seed The random seed.
+   */
+  sample(
+    value: number | SampleClauseNode,
+    method?: SampleMethod,
+    seed?: number
+  ): this {
     let clause;
     if (typeof value === 'number') {
       const perc = value > 0 && value < 1;
@@ -357,64 +352,58 @@ export class SelectQuery extends Query {
 
   /**
    * Add WHERE expressions.
-   * @param  {...FilterExpr} expr Expressions to add.
-   * @returns {this}
+   * @param expr Expressions to add.
    */
-  where(...expr) {
+  where(...expr: FilterExpr[]): this {
     this._where = this._where.concat(exprList(expr, asVerbatim));
     return this;
   }
 
   /**
    * Set WHERE expressions, replacing any prior expressions.
-   * @param {...FilterExpr} expr Expressions to add.
-   * @returns {this}
+   * @param expr Expressions to add.
    */
-  setWhere(...expr) {
+  setWhere(...expr: FilterExpr[]): this {
     this._where = [];
     return this.where(...expr);
   }
 
   /**
    * Add GROUP BY expressions.
-   * @param  {...GroupByExpr} expr Expressions to add.
-   * @returns {this}
+   * @param expr Expressions to add.
    */
-  groupby(...expr) {
-    this._groupby = this._groupby.concat(exprList(expr));
+  groupby(...expr: GroupByExpr[]): this {
+    this._groupby = this._groupby.concat(nodeList(expr));
     return this;
   }
 
   /**
    * Set GROUP BY expressions, replacing any prior expressions.
-   * @param {...GroupByExpr} expr Expressions to add.
-   * @returns {this}
+   * @param expr Expressions to add.
    */
-  setGroupby(...expr) {
+  setGroupby(...expr: GroupByExpr[]): this {
     this._groupby = [];
     return this.groupby(...expr);
   }
 
   /**
    * Add HAVING expressions.
-   * @param {...FilterExpr} expr Expressions to add.
-   * @returns {this}
+   * @param expr Expressions to add.
    */
-  having(...expr) {
+  having(...expr: FilterExpr[]): this {
     this._having = this._having.concat(exprList(expr, asVerbatim));
     return this;
   }
 
   /**
    * Add WINDOW definitions.
-   * @param  {...any} expr Expressions to add.
-   * @returns {this}
+   * @param expr Window definitions to add.
    */
-  window(...expr) {
-    const list = [];
+  window(...expr: (Record<string, WindowDefNode> | null)[]): this {
+    const list: WindowClauseNode[] = [];
     expr.flat().forEach(e => {
       if (e != null) for (const name in e) {
-        list.push(new WindowClauseNode(unquote(name), e[name]));
+        list.push(new WindowClauseNode(unquote(name)!, e[name]));
       }
     });
     this._window = this._window.concat(list);
@@ -423,22 +412,20 @@ export class SelectQuery extends Query {
 
   /**
    * Add QUALIFY expressions.
-   * @param {...FilterExpr} expr Expressions to add.
-   * @returns {this}
+   * @param expr Expressions to add.
    */
-  qualify(...expr) {
+  qualify(...expr: FilterExpr[]): this {
     this._qualify = this._qualify.concat(exprList(expr, asVerbatim));
     return this;
   }
 
   /**
    * Generate a SQL query string.
-   * @returns {string}
    */
   toString() {
     const {
       _with, _select, _distinct, _from, _sample, _where, _groupby,
-      _having, _window, _qualify, _orderby, _limit, _offset
+      _having, _window, _qualify, _orderby, _limitPerc, _limit, _offset
     } = this;
     const sql = [];
 
@@ -484,35 +471,37 @@ export class SelectQuery extends Query {
     if (_orderby.length) sql.push(`ORDER BY ${_orderby.join(', ')}`);
 
     // LIMIT
-    if (Number.isFinite(_limit)) sql.push(`LIMIT ${_limit}`);
+    if (_limit) sql.push(`LIMIT ${_limit}${_limitPerc ? '%' : ''}`);
 
     // OFFSET
-    if (Number.isFinite(_offset)) sql.push(`OFFSET ${_offset}`);
+    if (_offset) sql.push(`OFFSET ${_offset}`);
 
     return sql.join(' ');
   }
 }
 
 export class DescribeQuery extends SQLNode {
+  readonly query: Query;
+
   /**
    * Instantiate a describe query.
+   * @param query The query to describe.
    */
-  constructor(query) {
+  constructor(query: Query) {
     super(DESCRIBE_QUERY);
     this.query = query;
   }
 
   /**
    * Clone this describe query.
-   * @returns {DescribeQuery}
    */
-  clone() {
+  clone(): this {
+    // @ts-expect-error creates describe query
     return new DescribeQuery(this.query.clone());
   }
 
   /**
    * Generate a SQL query string.
-   * @returns {string}
    */
   toString() {
     return `DESCRIBE ${this.query}`;
@@ -520,51 +509,53 @@ export class DescribeQuery extends SQLNode {
 }
 
 export class SetOperation extends Query {
+  /** The set operation to perform. */
+  readonly op: string;
+  /** The input queries to the set operation. */
+  readonly queries: Query[];
+
   /**
    * Instantiate a new set operation instance.
-   * @param {string} op The set operation.
-   * @param {Query[]} queries The subqueries.
+   * @param op The set operation.
+   * @param queries The subqueries.
    */
-  constructor(op, queries) {
+  constructor(op: string, queries: Query[]) {
     super(SET_OPERATION);
-    /**
-     * @type {string}
-     * @readonly
-     */
     this.op = op;
-    /**
-     * @type {Query[]}
-     * @readonly
-     */
     this.queries = queries;
   }
 
+  /**
+   * Add a pointer to the query for which this query is a CTE.
+   * @param query The query for which this query is a CTE.
+   */
+  setCteFor(query: Query | null) {
+    super.setCteFor(query);
+    const { queries, cteFor } = this;
+    if (cteFor) queries.forEach(q => q.setCteFor(cteFor));
+  }
+
   /**
    * Return a list of subqueries.
-   * @returns {Query[]}
    */
   get subqueries() {
-    const { queries, cteFor } = this;
-    // TODO: revisit this?
-    if (cteFor) queries.forEach(q => q.cteFor = cteFor);
-    return queries;
+    return this.queries;
   }
 
   /**
    * Clone this set operation.
-   * @returns {SetOperation}
    */
-  clone() {
+  clone(): this {
     const { op, queries, ...rest } = this;
+    // @ts-expect-error creates set operation
     return Object.assign(new SetOperation(op, queries), rest);
   }
 
   /**
    * Generate a SQL query string.
-   * @returns {string}
    */
   toString() {
-    const { op, queries, _with, _orderby, _limit, _offset } = this;
+    const { op, queries, _with, _orderby, _limitPerc, _limit, _offset } = this;
     const sql = [];
 
     // WITH
@@ -577,75 +568,72 @@ export class SetOperation extends Query {
     if (_orderby.length) sql.push(`ORDER BY ${_orderby.join(', ')}`);
 
     // LIMIT
-    if (Number.isFinite(_limit)) sql.push(`LIMIT ${_limit}`);
+    if (_limit) sql.push(`LIMIT ${_limit}${_limitPerc ? '%' : ''}`);
 
     // OFFSET
-    if (Number.isFinite(_offset)) sql.push(`OFFSET ${_offset}`);
+    if (_offset) sql.push(`OFFSET ${_offset}`);
 
     return sql.join(' ');
   }
 }
 
 class WithClause {
+  /** The common table expressions (CTE). */
+  readonly _with: WithExpr[];
+
   /**
    * Instantiate a new WITH clause instance.
-   * @param {...WithExpr} expr The WITH CTE queries.
+   * @param expr The WITH CTE queries.
    */
-  constructor(...expr) {
+  constructor(...expr: WithExpr[]) {
     this._with = expr;
   }
 
   /**
    * Create a new select query with the given SELECT expressions.
-   * @param {...SelectExpr} expr The SELECT expressions.
-   * @returns {SelectQuery}
+   * @param expr The SELECT expressions.
    */
-  select(...expr) {
+  select(...expr: SelectExpr[]) {
     return Query.select(...expr).with(...this._with);
   }
 
   /**
    * Create a new select query with the given FROM expressions.
-   * @param {...FromExpr} expr The FROM expressions.
-   * @returns {SelectQuery}
+   * @param expr The FROM expressions.
    */
-  from(...expr) {
+  from(...expr: FromExpr[]) {
     return Query.from(...expr).with(...this._with);
   }
 
   /**
    * Create a new UNION set operation over the given queries.
-   * @param {...Query} queries The queries.
-   * @returns {SetOperation}
+   * @param queries The queries.
    */
-  union(...queries) {
+  union(...queries: Query[]) {
     return Query.union(...queries).with(...this._with);
   }
 
   /**
    * Create a new UNION ALL set operation over the given queries.
-   * @param {...Query} queries The queries.
-   * @returns {SetOperation}
+   * @param queries The queries.
    */
-  unionAll(...queries) {
+  unionAll(...queries: Query[]) {
     return Query.unionAll(...queries).with(...this._with);
   }
 
   /**
    * Create a new INTERSECT set operation over the given queries.
-   * @param {...Query} queries The queries.
-   * @returns {SetOperation}
+   * @param queries The queries.
    */
-  intersect(...queries) {
+  intersect(...queries: Query[]) {
     return Query.intersect(...queries).with(...this._with);
   }
 
   /**
    * Create a new EXCEPT set operation over the given queries.
-   * @param {...Query} queries The queries.
-   * @returns {SetOperation}
+   * @param queries The queries.
    */
-  except(...queries) {
+  except(...queries: Query[]) {
     return Query.except(...queries).with(...this._with);
   }
 }