drizzle-databend 0.1.11 → 0.1.13

package/src/databend-core/query-builders/select.ts

@@ -0,0 +1,412 @@
+import { entityKind, is } from 'drizzle-orm/entity';
+import { TypedQueryBuilder } from 'drizzle-orm/query-builders/query-builder';
+import { QueryPromise } from 'drizzle-orm/query-promise';
+import { SelectionProxyHandler } from 'drizzle-orm/selection-proxy';
+import { SQL, View } from 'drizzle-orm/sql/sql';
+import { Subquery } from 'drizzle-orm/subquery';
+import { Table } from 'drizzle-orm/table';
+// @ts-expect-error - tracer is exported at runtime but not in .d.ts
+import { tracer } from 'drizzle-orm/tracing';
+import {
+  // @ts-expect-error - exported at runtime but not in .d.ts
+  applyMixins,
+  getTableColumns,
+  // @ts-expect-error - exported at runtime but not in .d.ts
+  getTableLikeName,
+  haveSameKeys,
+  // @ts-expect-error - exported at runtime but not in .d.ts
+  orderSelectedFields,
+} from 'drizzle-orm/utils';
+import { ViewBaseConfig } from 'drizzle-orm/view-common';
+import { DatabendColumn } from '../columns/common.ts';
+import { DatabendViewBase } from '../view-base.ts';
+
+export class DatabendSelectBuilder {
+  static readonly [entityKind]: string = 'DatabendSelectBuilder';
+
+  fields: any;
+  session: any;
+  dialect: any;
+  withList: any[] = [];
+  distinct: any;
+
+  constructor(config: any) {
+    this.fields = config.fields;
+    this.session = config.session;
+    this.dialect = config.dialect;
+    if (config.withList) {
+      this.withList = config.withList;
+    }
+    this.distinct = config.distinct;
+  }
+
+  from(source: any): any {
+    const isPartialSelect = !!this.fields;
+
+    let fields: any;
+    if (this.fields) {
+      fields = this.fields;
+    } else if (is(source, Subquery)) {
+      fields = Object.fromEntries(
+        Object.keys(source._.selectedFields).map((key) => [key, (source as any)[key]])
+      );
+    } else if (is(source, DatabendViewBase)) {
+      fields = (source as any)[ViewBaseConfig].selectedFields;
+    } else if (is(source, SQL)) {
+      fields = {};
+    } else {
+      fields = getTableColumns(source);
+    }
+
+    return new DatabendSelectBase({
+      table: source,
+      fields,
+      isPartialSelect,
+      session: this.session,
+      dialect: this.dialect,
+      withList: this.withList,
+      distinct: this.distinct,
+    });
+  }
+}
+
+class DatabendSelectQueryBuilderBase extends TypedQueryBuilder<any, any> {
+  static readonly [entityKind]: string = 'DatabendSelectQueryBuilder';
+
+  _: any;
+  config: any;
+  joinsNotNullableMap: any;
+  tableName: any;
+  isPartialSelect: any;
+  session: any;
+  dialect: any;
+
+  constructor({ table, fields, isPartialSelect, session, dialect, withList, distinct }: any) {
+    super();
+    this.config = {
+      withList,
+      table,
+      fields: { ...fields },
+      distinct,
+      setOperators: [],
+    };
+    this.isPartialSelect = isPartialSelect;
+    this.session = session;
+    this.dialect = dialect;
+    this._ = {
+      selectedFields: fields,
+    };
+    this.tableName = getTableLikeName(table);
+    this.joinsNotNullableMap = typeof this.tableName === 'string' ? { [this.tableName]: true } : {};
+  }
+
+  private createJoin(joinType: string) {
+    return (table: any, on: any) => {
+      const baseTableName = this.tableName;
+      const tableName = getTableLikeName(table);
+
+      if (typeof tableName === 'string' && this.config.joins?.some((join: any) => join.alias === tableName)) {
+        throw new Error(`Alias "${tableName}" is already used in this query`);
+      }
+
+      if (!this.isPartialSelect) {
+        if (Object.keys(this.joinsNotNullableMap).length === 1 && typeof baseTableName === 'string') {
+          this.config.fields = {
+            [baseTableName]: this.config.fields,
+          };
+        }
+        if (typeof tableName === 'string' && !is(table, SQL)) {
+          const selection = is(table, Subquery)
+            ? table._.selectedFields
+            : is(table, View)
+              ? (table as any)[ViewBaseConfig].selectedFields
+              : table[(Table as any).Symbol.Columns];
+          this.config.fields[tableName] = selection;
+        }
+      }
+
+      if (typeof on === 'function') {
+        on = on(
+          new Proxy(
+            this.config.fields,
+            new SelectionProxyHandler({ sqlAliasedBehavior: 'sql', sqlBehavior: 'sql' })
+          )
+        );
+      }
+
+      if (!this.config.joins) {
+        this.config.joins = [];
+      }
+      this.config.joins.push({ on, table, joinType, alias: tableName });
+
+      if (typeof tableName === 'string') {
+        switch (joinType) {
+          case 'left':
+            this.joinsNotNullableMap[tableName] = false;
+            break;
+          case 'right':
+            this.joinsNotNullableMap = Object.fromEntries(
+              Object.entries(this.joinsNotNullableMap).map(([key]) => [key, false])
+            );
+            this.joinsNotNullableMap[tableName] = true;
+            break;
+          case 'inner':
+            this.joinsNotNullableMap[tableName] = true;
+            break;
+          case 'full':
+            this.joinsNotNullableMap = Object.fromEntries(
+              Object.entries(this.joinsNotNullableMap).map(([key]) => [key, false])
+            );
+            this.joinsNotNullableMap[tableName] = false;
+            break;
+        }
+      }
+      return this;
+    };
+  }
+
+  leftJoin = this.createJoin('left');
+  rightJoin = this.createJoin('right');
+  innerJoin = this.createJoin('inner');
+  fullJoin = this.createJoin('full');
+
+  crossJoin = (table: any) => {
+    const baseTableName = this.tableName;
+    const tableName = getTableLikeName(table);
+
+    if (typeof tableName === 'string' && this.config.joins?.some((join: any) => join.alias === tableName)) {
+      throw new Error(`Alias "${tableName}" is already used in this query`);
+    }
+
+    if (!this.isPartialSelect) {
+      if (Object.keys(this.joinsNotNullableMap).length === 1 && typeof baseTableName === 'string') {
+        this.config.fields = {
+          [baseTableName]: this.config.fields,
+        };
+      }
+      if (typeof tableName === 'string' && !is(table, SQL)) {
+        const selection = is(table, Subquery)
+          ? table._.selectedFields
+          : is(table, View)
+            ? (table as any)[ViewBaseConfig].selectedFields
+            : table[(Table as any).Symbol.Columns];
+        this.config.fields[tableName] = selection;
+      }
+    }
+
+    if (!this.config.joins) {
+      this.config.joins = [];
+    }
+    this.config.joins.push({ on: undefined, table, joinType: 'cross', alias: tableName });
+
+    if (typeof tableName === 'string') {
+      this.joinsNotNullableMap[tableName] = true;
+    }
+    return this;
+  };
+
+  private createSetOperator(type: string, isAll: boolean) {
+    return (rightSelection: any) => {
+      const rightSelect = typeof rightSelection === 'function'
+        ? rightSelection(getDatabendSetOperators())
+        : rightSelection;
+      if (!haveSameKeys(this.getSelectedFields(), rightSelect.getSelectedFields())) {
+        throw new Error(
+          'Set operator error (union / intersect / except): selected fields are not the same or are in a different order'
+        );
+      }
+      this.config.setOperators.push({ type, isAll, rightSelect });
+      return this;
+    };
+  }
+
+  union = this.createSetOperator('union', false);
+  unionAll = this.createSetOperator('union', true);
+  intersect = this.createSetOperator('intersect', false);
+  intersectAll = this.createSetOperator('intersect', true);
+  except = this.createSetOperator('except', false);
+  exceptAll = this.createSetOperator('except', true);
+
+  /** @internal */
+  addSetOperators(setOperators: any[]) {
+    this.config.setOperators.push(...setOperators);
+    return this;
+  }
+
+  where(where: any) {
+    if (typeof where === 'function') {
+      where = where(
+        new Proxy(
+          this.config.fields,
+          new SelectionProxyHandler({ sqlAliasedBehavior: 'sql', sqlBehavior: 'sql' })
+        )
+      );
+    }
+    this.config.where = where;
+    return this;
+  }
+
+  having(having: any) {
+    if (typeof having === 'function') {
+      having = having(
+        new Proxy(
+          this.config.fields,
+          new SelectionProxyHandler({ sqlAliasedBehavior: 'sql', sqlBehavior: 'sql' })
+        )
+      );
+    }
+    this.config.having = having;
+    return this;
+  }
+
+  groupBy(...columns: any[]) {
+    if (typeof columns[0] === 'function') {
+      const groupBy = columns[0](
+        new Proxy(
+          this.config.fields,
+          new SelectionProxyHandler({ sqlAliasedBehavior: 'alias', sqlBehavior: 'sql' })
+        )
+      );
+      this.config.groupBy = Array.isArray(groupBy) ? groupBy : [groupBy];
+    } else {
+      this.config.groupBy = columns;
+    }
+    return this;
+  }
+
+  orderBy(...columns: any[]) {
+    if (typeof columns[0] === 'function') {
+      const orderBy = columns[0](
+        new Proxy(
+          this.config.fields,
+          new SelectionProxyHandler({ sqlAliasedBehavior: 'alias', sqlBehavior: 'sql' })
+        )
+      );
+      const orderByArray = Array.isArray(orderBy) ? orderBy : [orderBy];
+      if (this.config.setOperators.length > 0) {
+        this.config.setOperators.at(-1).orderBy = orderByArray;
+      } else {
+        this.config.orderBy = orderByArray;
+      }
+    } else {
+      const orderByArray = columns;
+      if (this.config.setOperators.length > 0) {
+        this.config.setOperators.at(-1).orderBy = orderByArray;
+      } else {
+        this.config.orderBy = orderByArray;
+      }
+    }
+    return this;
+  }
+
+  limit(limit: any) {
+    if (this.config.setOperators.length > 0) {
+      this.config.setOperators.at(-1).limit = limit;
+    } else {
+      this.config.limit = limit;
+    }
+    return this;
+  }
+
+  offset(offset: any) {
+    if (this.config.setOperators.length > 0) {
+      this.config.setOperators.at(-1).offset = offset;
+    } else {
+      this.config.offset = offset;
+    }
+    return this;
+  }
+
+  /** @internal */
+  getSQL() {
+    return this.dialect.buildSelectQuery(this.config);
+  }
+
+  toSQL() {
+    const { typings: _typings, ...rest } = this.dialect.sqlToQuery(this.getSQL());
+    return rest;
+  }
+
+  as(alias: string) {
+    return new Proxy(
+      new Subquery(this.getSQL(), this.config.fields, alias),
+      new SelectionProxyHandler({ alias, sqlAliasedBehavior: 'alias', sqlBehavior: 'error' })
+    );
+  }
+
+  /** @internal */
+  getSelectedFields() {
+    return new Proxy(
+      this.config.fields,
+      new SelectionProxyHandler({ alias: this.tableName, sqlAliasedBehavior: 'alias', sqlBehavior: 'error' })
+    );
+  }
+
+  $dynamic() {
+    return this;
+  }
+}
+
+export class DatabendSelectBase extends DatabendSelectQueryBuilderBase {
+  static override readonly [entityKind]: string = 'DatabendSelect';
+
+  /** @internal */
+  _prepare(name?: string) {
+    const { session, config, dialect, joinsNotNullableMap } = this;
+    if (!session) {
+      throw new Error('Cannot execute a query on a query builder. Please use a database instance instead.');
+    }
+    return tracer.startActiveSpan('drizzle.prepareQuery', () => {
+      const fieldsList = orderSelectedFields(config.fields);
+      const query = session.prepareQuery(dialect.sqlToQuery(this.getSQL()), fieldsList, name, true);
+      query.joinsNotNullableMap = joinsNotNullableMap;
+      return query;
+    });
+  }
+
+  prepare(name: string) {
+    return this._prepare(name);
+  }
+
+  execute = (placeholderValues?: any) => {
+    return tracer.startActiveSpan('drizzle.operation', () => {
+      return this._prepare().execute(placeholderValues);
+    });
+  };
+}
+
+applyMixins(DatabendSelectBase, [QueryPromise]);
+
+function createSetOperator(type: string, isAll: boolean) {
+  return (leftSelect: any, rightSelect: any, ...restSelects: any[]) => {
+    const setOperators = [rightSelect, ...restSelects].map((select) => ({
+      type,
+      isAll,
+      rightSelect: select,
+    }));
+    for (const setOperator of setOperators) {
+      if (!haveSameKeys(leftSelect.getSelectedFields(), setOperator.rightSelect.getSelectedFields())) {
+        throw new Error(
+          'Set operator error (union / intersect / except): selected fields are not the same or are in a different order'
+        );
+      }
+    }
+    return leftSelect.addSetOperators(setOperators);
+  };
+}
+
+const getDatabendSetOperators = () => ({
+  union,
+  unionAll,
+  intersect,
+  intersectAll,
+  except,
+  exceptAll,
+});
+
+export const union = createSetOperator('union', false);
+export const unionAll = createSetOperator('union', true);
+export const intersect = createSetOperator('intersect', false);
+export const intersectAll = createSetOperator('intersect', true);
+export const except = createSetOperator('except', false);
+export const exceptAll = createSetOperator('except', true);

package/src/databend-core/query-builders/update.ts

@@ -0,0 +1,82 @@
+import { entityKind, is } from 'drizzle-orm/entity';
+import { QueryPromise } from 'drizzle-orm/query-promise';
+import { SQL } from 'drizzle-orm/sql/sql';
+import { Table } from 'drizzle-orm/table';
+// @ts-expect-error - tracer is exported at runtime but not in .d.ts
+import { tracer } from 'drizzle-orm/tracing';
+// @ts-expect-error - mapUpdateSet is exported at runtime but not in .d.ts
+import { mapUpdateSet } from 'drizzle-orm/utils';
+
+export class DatabendUpdateBuilder {
+  static readonly [entityKind]: string = 'DatabendUpdateBuilder';
+
+  constructor(
+    private table: any,
+    private session: any,
+    private dialect: any,
+    private withList?: any[]
+  ) {}
+
+  set(values: any) {
+    return new DatabendUpdateBase(
+      this.table,
+      mapUpdateSet(this.table, values),
+      this.session,
+      this.dialect,
+      this.withList
+    );
+  }
+}
+
+export class DatabendUpdateBase extends QueryPromise<any> {
+  static readonly [entityKind]: string = 'DatabendUpdate';
+
+  config: any;
+
+  constructor(
+    table: any,
+    set: any,
+    private session: any,
+    private dialect: any,
+    withList?: any[]
+  ) {
+    super();
+    this.config = { set, table, withList };
+  }
+
+  where(where: any) {
+    this.config.where = where;
+    return this;
+  }
+
+  /** @internal */
+  getSQL() {
+    return this.dialect.buildUpdateQuery(this.config);
+  }
+
+  toSQL() {
+    const { typings: _typings, ...rest } = this.dialect.sqlToQuery(this.getSQL());
+    return rest;
+  }
+
+  /** @internal */
+  _prepare(name?: string) {
+    return tracer.startActiveSpan('drizzle.prepareQuery', () => {
+      return this.session.prepareQuery(this.dialect.sqlToQuery(this.getSQL()), undefined, name, true);
+    });
+  }
+
+  prepare(name: string) {
+    return this._prepare(name);
+  }
+
+  execute = (placeholderValues?: any) => {
+    return tracer.startActiveSpan('drizzle.operation', () => {
+      return this._prepare().execute(placeholderValues);
+    });
+  };
+
+  $dynamic() {
+    return this;
+  }
+}

package/src/databend-core/schema.ts

@@ -0,0 +1,29 @@
+import { entityKind, is } from 'drizzle-orm/entity';
+import { SQL, sql } from 'drizzle-orm/sql/sql';
+import { databendTableWithSchema } from './table.ts';
+
+export class DatabendSchema {
+  static readonly [entityKind]: string = 'DatabendSchema';
+
+  constructor(public schemaName: string) {}
+
+  table = (name: string, columns: any, extraConfig?: any): any => {
+    return databendTableWithSchema(name, columns, extraConfig, this.schemaName);
+  };
+
+  getSQL(): SQL {
+    return new SQL([sql.identifier(this.schemaName)]);
+  }
+
+  shouldOmitSQLParens() {
+    return true;
+  }
+}
+
+export function isDatabendSchema(obj: unknown): obj is DatabendSchema {
+  return is(obj, DatabendSchema);
+}
+
+export function databendSchema(name: string) {
+  return new DatabendSchema(name);
+}

package/src/databend-core/session.ts

@@ -0,0 +1,85 @@
+import { entityKind } from 'drizzle-orm/entity';
+import { TransactionRollbackError } from 'drizzle-orm/errors';
+// @ts-expect-error - tracer is exported at runtime but not in .d.ts
+import { tracer } from 'drizzle-orm/tracing';
+import { DatabendDatabase } from './db.ts';
+
+export class DatabendPreparedQuery {
+  static readonly [entityKind]: string = 'DatabendPreparedQuery';
+
+  /** @internal */
+  joinsNotNullableMap: Record<string, boolean> | undefined;
+
+  constructor(public query: any) {}
+
+  getQuery() {
+    return this.query;
+  }
+
+  mapResult(response: any, _isFromBatch?: boolean) {
+    return response;
+  }
+}
+
+export class DatabendSession {
+  static readonly [entityKind]: string = 'DatabendSession';
+
+  constructor(public dialect: any) {}
+
+  execute(query: any): any {
+    return tracer.startActiveSpan('drizzle.operation', () => {
+      const prepared = tracer.startActiveSpan('drizzle.prepareQuery', () => {
+        return this.prepareQuery(
+          this.dialect.sqlToQuery(query),
+          undefined,
+          undefined,
+          false
+        );
+      });
+      return prepared.execute(undefined);
+    });
+  }
+
+  all(query: any): any {
+    return this.prepareQuery(
+      this.dialect.sqlToQuery(query),
+      undefined,
+      undefined,
+      false
+    ).all();
+  }
+
+  async count(sql: any): Promise<number> {
+    const res = await this.execute(sql);
+    return Number(res[0]['count']);
+  }
+
+  prepareQuery(
+    _query: any,
+    _fields: any,
+    _name: any,
+    _isResponseInArrayMode: boolean,
+    _customResultMapper?: any
+  ): any {
+    throw new Error('prepareQuery not implemented');
+  }
+
+  transaction(_transaction: any, _config?: any): Promise<any> {
+    throw new Error('transaction not implemented');
+  }
+}
+
+export class DatabendTransaction extends DatabendDatabase {
+  static override readonly [entityKind]: string = 'DatabendTransaction';
+
+  declare readonly schema: any;
+
+  constructor(dialect: any, session: any, schema: any, public nestedIndex = 0) {
+    super(dialect, session, schema);
+    this.schema = schema;
+  }
+
+  rollback(): never {
+    throw new TransactionRollbackError();
+  }
+}

package/src/databend-core/subquery.ts

@@ -0,0 +1 @@
+// Databend subquery types - re-exported from drizzle-orm core

package/src/databend-core/table.ts

@@ -0,0 +1,67 @@
+import { entityKind } from 'drizzle-orm/entity';
+import { Table } from 'drizzle-orm/table';
+import { getDatabendColumnBuilders } from './columns/all.ts';
+
+export class DatabendTable extends Table {
+  static readonly [entityKind]: string = 'DatabendTable';
+
+  /** @internal */
+  static Symbol = Object.assign({}, (Table as any).Symbol, {});
+
+  constructor(name: string, schema: string | undefined, baseName: string) {
+    super(name, schema, baseName);
+    (this as any)[(Table as any).Symbol.ExtraConfigBuilder] = undefined;
+    (this as any)[(Table as any).Symbol.ExtraConfigColumns] = {};
+  }
+}
+
+function databendTableWithSchema(
+  name: string,
+  columns: any,
+  extraConfig: any,
+  schema: string | undefined,
+  baseName = name
+): any {
+  const rawTable = new DatabendTable(name, schema, baseName);
+  const parsedColumns = typeof columns === 'function' ? columns(getDatabendColumnBuilders()) : columns;
+
+  const builtColumns = Object.fromEntries(
+    Object.entries(parsedColumns).map(([name, colBuilderBase]: [string, any]) => {
+      const colBuilder = colBuilderBase;
+      colBuilder.setName(name);
+      const column = colBuilder.build(rawTable);
+      return [name, column];
+    })
+  );
+
+  const builtColumnsForExtraConfig = Object.fromEntries(
+    Object.entries(parsedColumns).map(([name, colBuilderBase]: [string, any]) => {
+      const colBuilder = colBuilderBase;
+      colBuilder.setName(name);
+      const column = colBuilder.buildExtraConfigColumn(rawTable);
+      return [name, column];
+    })
+  );
+
+  const table = Object.assign(rawTable, builtColumns);
+  table[(Table as any).Symbol.Columns] = builtColumns;
+  table[(Table as any).Symbol.ExtraConfigColumns] = builtColumnsForExtraConfig;
+
+  if (extraConfig) {
+    table[(DatabendTable as any).Symbol.ExtraConfigBuilder] = extraConfig;
+  }
+
+  return table;
+}
+
+export const databendTable = (name: string, columns: any, extraConfig?: any): any => {
+  return databendTableWithSchema(name, columns, extraConfig, undefined);
+};
+
+export function databendTableCreator(customizeTableName: (name: string) => string) {
+  return (name: string, columns: any, extraConfig?: any) => {
+    return databendTableWithSchema(customizeTableName(name), columns, extraConfig, undefined, name);
+  };
+}
+
+export { databendTableWithSchema };