drizzle-databend 0.1.0

package/src/session.ts

@@ -0,0 +1,311 @@
+import { entityKind } from 'drizzle-orm/entity';
+import type { Logger } from 'drizzle-orm/logger';
+import { NoopLogger } from 'drizzle-orm/logger';
+import { PgTransaction } from 'drizzle-orm/pg-core';
+import type { SelectedFieldsOrdered } from 'drizzle-orm/pg-core/query-builders/select.types';
+import type {
+  PgTransactionConfig,
+  PreparedQueryConfig,
+  PgQueryResultHKT,
+} from 'drizzle-orm/pg-core/session';
+import { PgPreparedQuery, PgSession } from 'drizzle-orm/pg-core/session';
+import type {
+  RelationalSchemaConfig,
+  TablesRelationalConfig,
+} from 'drizzle-orm/relations';
+import { fillPlaceholders, type Query, SQL, sql } from 'drizzle-orm/sql/sql';
+import type { Assume } from 'drizzle-orm/utils';
+import { mapResultRow } from './sql/result-mapper.ts';
+import { TransactionRollbackError } from 'drizzle-orm/errors';
+import type { DatabendDialect } from './dialect.ts';
+import type {
+  DatabendClientLike,
+  DatabendConnectionPool,
+  RowData,
+} from './client.ts';
+import {
+  executeArraysOnClient,
+  executeOnClient,
+  prepareParams,
+  isPool,
+} from './client.ts';
+import type { Connection } from 'databend-driver';
+
+export type { DatabendClientLike, RowData } from './client.ts';
+
+export class DatabendPreparedQuery<
+  T extends PreparedQueryConfig,
+> extends PgPreparedQuery<T> {
+  static readonly [entityKind]: string = 'DatabendPreparedQuery';
+
+  constructor(
+    private client: DatabendClientLike,
+    private queryString: string,
+    private params: unknown[],
+    private logger: Logger,
+    private fields: SelectedFieldsOrdered | undefined,
+    private _isResponseInArrayMode: boolean,
+    private customResultMapper:
+      | ((rows: unknown[][]) => T['execute'])
+      | undefined
+  ) {
+    super({ sql: queryString, params });
+  }
+
+  async execute(
+    placeholderValues: Record<string, unknown> | undefined = {}
+  ): Promise<T['execute']> {
+    const params = prepareParams(
+      fillPlaceholders(this.params, placeholderValues)
+    );
+    this.logger.logQuery(this.queryString, params);
+
+    const { fields, joinsNotNullableMap, customResultMapper } =
+      this as typeof this & { joinsNotNullableMap?: Record<string, boolean> };
+
+    if (fields) {
+      const { rows } = await executeArraysOnClient(
+        this.client,
+        this.queryString,
+        params
+      );
+
+      if (rows.length === 0) {
+        return [] as T['execute'];
+      }
+
+      return customResultMapper
+        ? customResultMapper(rows)
+        : rows.map((row) =>
+            mapResultRow<T['execute']>(fields, row, joinsNotNullableMap)
+          );
+    }
+
+    const rows = await executeOnClient(this.client, this.queryString, params);
+
+    return rows as T['execute'];
+  }
+
+  all(
+    placeholderValues: Record<string, unknown> | undefined = {}
+  ): Promise<T['all']> {
+    return this.execute(placeholderValues);
+  }
+
+  isResponseInArrayMode(): boolean {
+    return this._isResponseInArrayMode;
+  }
+}
+
+export interface DatabendSessionOptions {
+  logger?: Logger;
+}
+
+export class DatabendSession<
+  TFullSchema extends Record<string, unknown> = Record<string, never>,
+  TSchema extends TablesRelationalConfig = Record<string, never>,
+> extends PgSession<DatabendQueryResultHKT, TFullSchema, TSchema> {
+  static readonly [entityKind]: string = 'DatabendSession';
+
+  protected override dialect: DatabendDialect;
+  private logger: Logger;
+  private rollbackOnly = false;
+
+  constructor(
+    private client: DatabendClientLike,
+    dialect: DatabendDialect,
+    private schema: RelationalSchemaConfig<TSchema> | undefined,
+    private options: DatabendSessionOptions = {}
+  ) {
+    super(dialect);
+    this.dialect = dialect;
+    this.logger = options.logger ?? new NoopLogger();
+  }
+
+  prepareQuery<T extends PreparedQueryConfig = PreparedQueryConfig>(
+    query: Query,
+    fields: SelectedFieldsOrdered | undefined,
+    name: string | undefined,
+    isResponseInArrayMode: boolean,
+    customResultMapper?: (rows: unknown[][]) => T['execute']
+  ): PgPreparedQuery<T> {
+    void name;
+    return new DatabendPreparedQuery(
+      this.client,
+      query.sql,
+      query.params,
+      this.logger,
+      fields,
+      isResponseInArrayMode,
+      customResultMapper
+    );
+  }
+
+  override async transaction<T>(
+    transaction: (tx: DatabendTransaction<TFullSchema, TSchema>) => Promise<T>,
+    config?: PgTransactionConfig
+  ): Promise<T> {
+    let pinnedConnection: Connection | undefined;
+    let pool: DatabendConnectionPool | undefined;
+
+    let clientForTx: DatabendClientLike = this.client;
+    if (isPool(this.client)) {
+      pool = this.client;
+      pinnedConnection = await pool.acquire();
+      clientForTx = pinnedConnection;
+    }
+
+    const session = new DatabendSession(
+      clientForTx,
+      this.dialect,
+      this.schema,
+      this.options
+    );
+
+    const tx = new DatabendTransaction<TFullSchema, TSchema>(
+      this.dialect,
+      session,
+      this.schema
+    );
+
+    try {
+      await tx.execute(sql`BEGIN`);
+
+      if (config) {
+        await tx.setTransaction(config);
+      }
+
+      try {
+        const result = await transaction(tx);
+        if (session.isRollbackOnly()) {
+          await tx.execute(sql`ROLLBACK`);
+          throw new TransactionRollbackError();
+        }
+        await tx.execute(sql`COMMIT`);
+        return result;
+      } catch (error) {
+        await tx.execute(sql`ROLLBACK`);
+        throw error;
+      }
+    } finally {
+      if (pinnedConnection && pool) {
+        await pool.release(pinnedConnection);
+      }
+    }
+  }
+
+  markRollbackOnly(): void {
+    this.rollbackOnly = true;
+  }
+
+  isRollbackOnly(): boolean {
+    return this.rollbackOnly;
+  }
+}
+
+type DatabendTransactionInternals<
+  TFullSchema extends Record<string, unknown> = Record<string, never>,
+  TSchema extends TablesRelationalConfig = Record<string, never>,
+> = {
+  dialect: DatabendDialect;
+  session: DatabendSession<TFullSchema, TSchema>;
+};
+
+type DatabendTransactionWithInternals<
+  TFullSchema extends Record<string, unknown> = Record<string, never>,
+  TSchema extends TablesRelationalConfig = Record<string, never>,
+> = DatabendTransactionInternals<TFullSchema, TSchema> &
+  DatabendTransaction<TFullSchema, TSchema>;
+
+const VALID_TRANSACTION_ISOLATION_LEVELS = new Set<string>([
+  'read uncommitted',
+  'read committed',
+  'repeatable read',
+  'serializable',
+]);
+
+const VALID_TRANSACTION_ACCESS_MODES = new Set<string>([
+  'read only',
+  'read write',
+]);
+
+export class DatabendTransaction<
+  TFullSchema extends Record<string, unknown>,
+  TSchema extends TablesRelationalConfig,
+> extends PgTransaction<DatabendQueryResultHKT, TFullSchema, TSchema> {
+  static readonly [entityKind]: string = 'DatabendTransaction';
+
+  rollback(): never {
+    throw new TransactionRollbackError();
+  }
+
+  getTransactionConfigSQL(config: PgTransactionConfig): SQL {
+    if (
+      config.isolationLevel &&
+      !VALID_TRANSACTION_ISOLATION_LEVELS.has(config.isolationLevel)
+    ) {
+      throw new Error(
+        `Invalid transaction isolation level "${config.isolationLevel}". Expected one of: ${Array.from(
+          VALID_TRANSACTION_ISOLATION_LEVELS
+        ).join(', ')}.`
+      );
+    }
+
+    if (
+      config.accessMode &&
+      !VALID_TRANSACTION_ACCESS_MODES.has(config.accessMode)
+    ) {
+      throw new Error(
+        `Invalid transaction access mode "${config.accessMode}". Expected one of: ${Array.from(
+          VALID_TRANSACTION_ACCESS_MODES
+        ).join(', ')}.`
+      );
+    }
+
+    const chunks: string[] = [];
+    if (config.isolationLevel) {
+      chunks.push(`isolation level ${config.isolationLevel}`);
+    }
+    if (config.accessMode) {
+      chunks.push(config.accessMode);
+    }
+    return sql.raw(chunks.join(' '));
+  }
+
+  setTransaction(config: PgTransactionConfig): Promise<void> {
+    type Tx = DatabendTransactionWithInternals<TFullSchema, TSchema>;
+    return (this as unknown as Tx).session.execute(
+      sql`SET TRANSACTION ${this.getTransactionConfigSQL(config)}`
+    );
+  }
+
+  override async transaction<T>(
+    transaction: (tx: DatabendTransaction<TFullSchema, TSchema>) => Promise<T>
+  ): Promise<T> {
+    // Databend does not support savepoints. Use rollback-only fallback.
+    type Tx = DatabendTransactionWithInternals<TFullSchema, TSchema>;
+    const internals = this as unknown as Tx;
+
+    const nestedTx = new DatabendTransaction<TFullSchema, TSchema>(
+      internals.dialect,
+      internals.session,
+      this.schema,
+      this.nestedIndex + 1
+    );
+
+    return transaction(nestedTx).catch((error) => {
+      (
+        internals.session as DatabendSession<TFullSchema, TSchema>
+      ).markRollbackOnly();
+      throw error;
+    });
+  }
+}
+
+export type GenericRowData<T extends RowData = RowData> = T;
+
+export type GenericTableData<T = RowData> = T[];
+
+export interface DatabendQueryResultHKT extends PgQueryResultHKT {
+  type: GenericTableData<Assume<this['row'], RowData>>;
+}

package/src/sql/result-mapper.ts

@@ -0,0 +1,234 @@
+import {
+  Column,
+  SQL,
+  getTableName,
+  is,
+  type AnyColumn,
+  type DriverValueDecoder,
+  type SelectedFieldsOrdered,
+} from 'drizzle-orm';
+import {
+  PgCustomColumn,
+  PgDate,
+  PgDateString,
+  PgTime,
+  PgTimestamp,
+  PgTimestampString,
+} from 'drizzle-orm/pg-core';
+
+type SQLInternal<T = unknown> = SQL<T> & {
+  decoder: DriverValueDecoder<T, any>;
+};
+
+type DecoderInput<TDecoder extends DriverValueDecoder<unknown, unknown>> =
+  Parameters<TDecoder['mapFromDriverValue']>[0];
+
+function toDecoderInput<TDecoder extends DriverValueDecoder<unknown, unknown>>(
+  decoder: TDecoder,
+  value: unknown
+): DecoderInput<TDecoder> {
+  void decoder;
+  return value as DecoderInput<TDecoder>;
+}
+
+export function normalizeTimestampString(
+  value: unknown,
+  withTimezone: boolean
+): string | unknown {
+  if (value instanceof Date) {
+    const iso = value.toISOString().replace('T', ' ');
+    return withTimezone ? iso.replace('Z', '+00') : iso.replace('Z', '');
+  }
+  if (typeof value === 'string') {
+    const normalized = value.replace('T', ' ');
+    if (withTimezone) {
+      return normalized.includes('+') ? normalized : `${normalized}+00`;
+    }
+    return normalized.replace(/\+00$/, '');
+  }
+  return value;
+}
+
+export function normalizeTimestamp(
+  value: unknown,
+  withTimezone: boolean
+): Date | unknown {
+  if (value instanceof Date) {
+    return value;
+  }
+  if (typeof value === 'string') {
+    const hasOffset =
+      value.endsWith('Z') || /[+-]\d{2}:?\d{2}$/.test(value.trim());
+    const spaced = value.replace(' ', 'T');
+    const normalized = withTimezone || hasOffset ? spaced : `${spaced}+00`;
+    return new Date(normalized);
+  }
+  return value;
+}
+
+export function normalizeDateString(value: unknown): string | unknown {
+  if (value instanceof Date) {
+    return value.toISOString().slice(0, 10);
+  }
+  if (typeof value === 'string') {
+    return value.slice(0, 10);
+  }
+  return value;
+}
+
+export function normalizeDateValue(value: unknown): Date | unknown {
+  if (value instanceof Date) {
+    return value;
+  }
+  if (typeof value === 'string') {
+    return new Date(`${value.slice(0, 10)}T00:00:00Z`);
+  }
+  return value;
+}
+
+export function normalizeTime(value: unknown): string | unknown {
+  if (typeof value === 'bigint') {
+    const totalMillis = Number(value) / 1000;
+    const date = new Date(totalMillis);
+    return date.toISOString().split('T')[1]!.replace('Z', '');
+  }
+  if (value instanceof Date) {
+    return value.toISOString().split('T')[1]!.replace('Z', '');
+  }
+  return value;
+}
+
+function mapDriverValue(
+  decoder: DriverValueDecoder<unknown, unknown>,
+  rawValue: unknown
+): unknown {
+  if (is(decoder, PgTimestampString)) {
+    return decoder.mapFromDriverValue(
+      toDecoderInput(
+        decoder,
+        normalizeTimestampString(rawValue, decoder.withTimezone)
+      )
+    );
+  }
+
+  if (is(decoder, PgTimestamp)) {
+    const normalized = normalizeTimestamp(rawValue, decoder.withTimezone);
+    if (normalized instanceof Date) {
+      return normalized;
+    }
+    return decoder.mapFromDriverValue(toDecoderInput(decoder, normalized));
+  }
+
+  if (is(decoder, PgDateString)) {
+    return decoder.mapFromDriverValue(
+      toDecoderInput(decoder, normalizeDateString(rawValue))
+    );
+  }
+
+  if (is(decoder, PgDate)) {
+    return decoder.mapFromDriverValue(
+      toDecoderInput(decoder, normalizeDateValue(rawValue))
+    );
+  }
+
+  if (is(decoder, PgTime)) {
+    return decoder.mapFromDriverValue(
+      toDecoderInput(decoder, normalizeTime(rawValue))
+    );
+  }
+
+  return decoder.mapFromDriverValue(toDecoderInput(decoder, rawValue));
+}
+
+export function mapResultRow<TResult>(
+  columns: SelectedFieldsOrdered<AnyColumn>,
+  row: unknown[],
+  joinsNotNullableMap: Record<string, boolean> | undefined
+): TResult {
+  const nullifyMap: Record<string, string | false> = {};
+
+  const result = columns.reduce<Record<string, any>>(
+    (acc, { path, field }, columnIndex) => {
+      let decoder: DriverValueDecoder<unknown, unknown>;
+      if (is(field, Column)) {
+        decoder = field;
+      } else if (is(field, SQL)) {
+        decoder = (field as SQLInternal).decoder;
+      } else {
+        const col = field.sql.queryChunks.find((chunk) => is(chunk, Column));
+
+        if (is(col, PgCustomColumn)) {
+          decoder = col;
+        } else {
+          decoder = (field.sql as SQLInternal).decoder;
+        }
+      }
+      let node = acc;
+      for (const [pathChunkIndex, pathChunk] of path.entries()) {
+        if (pathChunkIndex < path.length - 1) {
+          if (!(pathChunk in node)) {
+            node[pathChunk] = {};
+          }
+          node = node[pathChunk];
+          continue;
+        }
+
+        const rawValue = row[columnIndex]!;
+
+        const value = (node[pathChunk] =
+          rawValue === null ? null : mapDriverValue(decoder, rawValue));
+
+        if (joinsNotNullableMap && is(field, Column) && path.length === 2) {
+          const objectName = path[0]!;
+          if (!(objectName in nullifyMap)) {
+            nullifyMap[objectName] =
+              value === null ? getTableName(field.table) : false;
+          } else if (
+            typeof nullifyMap[objectName] === 'string' &&
+            nullifyMap[objectName] !== getTableName(field.table)
+          ) {
+            nullifyMap[objectName] = false;
+          }
+          continue;
+        }
+
+        if (
+          joinsNotNullableMap &&
+          is(field, SQL.Aliased) &&
+          path.length === 2
+        ) {
+          const col = field.sql.queryChunks.find((chunk) => is(chunk, Column));
+          const tableName = col?.table && getTableName(col?.table);
+
+          if (!tableName) {
+            continue;
+          }
+
+          const objectName = path[0]!;
+
+          if (!(objectName in nullifyMap)) {
+            nullifyMap[objectName] = value === null ? tableName : false;
+            continue;
+          }
+
+          if (nullifyMap[objectName] && nullifyMap[objectName] !== tableName) {
+            nullifyMap[objectName] = false;
+          }
+          continue;
+        }
+      }
+      return acc;
+    },
+    {}
+  );
+
+  if (joinsNotNullableMap && Object.keys(nullifyMap).length > 0) {
+    for (const [objectName, tableName] of Object.entries(nullifyMap)) {
+      if (typeof tableName === 'string' && !joinsNotNullableMap[tableName]) {
+        result[objectName] = null;
+      }
+    }
+  }
+
+  return result as TResult;
+}

package/src/sql/selection.ts

@@ -0,0 +1,60 @@
+import { Column, SQL, getTableName, is, sql } from 'drizzle-orm';
+import type { SelectedFields } from 'drizzle-orm/pg-core';
+
+function mapEntries(
+  obj: Record<string, unknown>,
+  prefix?: string,
+  fullJoin = false
+): Record<string, unknown> {
+  return Object.fromEntries(
+    Object.entries(obj)
+      .filter(([key]) => key !== 'enableRLS')
+      .map(([key, value]) => {
+        const qualified = prefix ? `${prefix}.${key}` : key;
+
+        if (fullJoin && is(value, Column)) {
+          return [
+            key,
+            sql`${value}`
+              .mapWith(value)
+              .as(`${getTableName(value.table)}.${value.name}`),
+          ];
+        }
+
+        if (fullJoin && is(value, SQL)) {
+          const col = value
+            .getSQL()
+            .queryChunks.find((chunk) => is(chunk, Column));
+
+          const tableName = col?.table && getTableName(col?.table);
+
+          return [key, value.as(tableName ? `${tableName}.${key}` : key)];
+        }
+
+        if (is(value, SQL) || is(value, Column)) {
+          const aliased = is(value, SQL) ? value : sql`${value}`.mapWith(value);
+          return [key, aliased.as(qualified)];
+        }
+
+        if (is(value, SQL.Aliased)) {
+          return [key, value];
+        }
+
+        if (typeof value === 'object' && value !== null) {
+          return [
+            key,
+            mapEntries(value as Record<string, unknown>, qualified, fullJoin),
+          ];
+        }
+
+        return [key, value];
+      })
+  );
+}
+
+export function aliasFields(
+  fields: SelectedFields,
+  fullJoin = false
+): SelectedFields {
+  return mapEntries(fields, undefined, fullJoin) as SelectedFields;
+}