drizzle-databend 0.1.0

package/src/driver.ts

@@ -0,0 +1,268 @@
+import { Client } from 'databend-driver';
+import { entityKind } from 'drizzle-orm/entity';
+import type { Logger } from 'drizzle-orm/logger';
+import { DefaultLogger } from 'drizzle-orm/logger';
+import { PgDatabase } from 'drizzle-orm/pg-core/db';
+import {
+  createTableRelationsHelpers,
+  extractTablesRelationalConfig,
+  type ExtractTablesWithRelations,
+  type RelationalSchemaConfig,
+  type TablesRelationalConfig,
+} from 'drizzle-orm/relations';
+import { type DrizzleConfig } from 'drizzle-orm/utils';
+import type {
+  DatabendClientLike,
+  DatabendQueryResultHKT,
+  DatabendTransaction,
+} from './session.ts';
+import { DatabendSession } from './session.ts';
+import { DatabendDialect } from './dialect.ts';
+import { isPool, closeClientConnection } from './client.ts';
+import {
+  createDatabendConnectionPool,
+  type DatabendPoolConfig,
+} from './pool.ts';
+import type { Connection } from 'databend-driver';
+
+export interface DatabendDriverOptions {
+  logger?: Logger;
+}
+
+export class DatabendDriver {
+  static readonly [entityKind]: string = 'DatabendDriver';
+
+  constructor(
+    private client: DatabendClientLike,
+    private dialect: DatabendDialect,
+    private options: DatabendDriverOptions = {}
+  ) {}
+
+  createSession(
+    schema: RelationalSchemaConfig<TablesRelationalConfig> | undefined
+  ): DatabendSession<Record<string, unknown>, TablesRelationalConfig> {
+    return new DatabendSession(this.client, this.dialect, schema, {
+      logger: this.options.logger,
+    });
+  }
+}
+
+export interface DatabendDrizzleConfig<
+  TSchema extends Record<string, unknown> = Record<string, never>,
+> extends DrizzleConfig<TSchema> {
+  /** Pool configuration. Use size config or false to disable. */
+  pool?: DatabendPoolConfig | false;
+}
+
+export interface DatabendDrizzleConfigWithConnection<
+  TSchema extends Record<string, unknown> = Record<string, never>,
+> extends DatabendDrizzleConfig<TSchema> {
+  /** Connection DSN string */
+  connection: string;
+}
+
+export interface DatabendDrizzleConfigWithClient<
+  TSchema extends Record<string, unknown> = Record<string, never>,
+> extends DatabendDrizzleConfig<TSchema> {
+  /** Explicit client (connection or pool) */
+  client: DatabendClientLike;
+}
+
+function isConfigObject(data: unknown): data is Record<string, unknown> {
+  if (typeof data !== 'object' || data === null) return false;
+  if (data.constructor?.name !== 'Object') return false;
+  return (
+    'connection' in data ||
+    'client' in data ||
+    'pool' in data ||
+    'schema' in data ||
+    'logger' in data
+  );
+}
+
+function createFromClient<
+  TSchema extends Record<string, unknown> = Record<string, never>,
+>(
+  client: DatabendClientLike,
+  config: DatabendDrizzleConfig<TSchema> = {},
+  databendClient?: Client
+): DatabendDatabase<TSchema, ExtractTablesWithRelations<TSchema>> {
+  const dialect = new DatabendDialect();
+
+  const logger =
+    config.logger === true ? new DefaultLogger() : config.logger || undefined;
+
+  let schema: RelationalSchemaConfig<TablesRelationalConfig> | undefined;
+
+  if (config.schema) {
+    const tablesConfig = extractTablesRelationalConfig(
+      config.schema,
+      createTableRelationsHelpers
+    );
+    schema = {
+      fullSchema: config.schema,
+      schema: tablesConfig.tables,
+      tableNamesMap: tablesConfig.tableNamesMap,
+    };
+  }
+
+  const driver = new DatabendDriver(client, dialect, { logger });
+  const session = driver.createSession(schema);
+
+  const db = new DatabendDatabase(
+    dialect,
+    session,
+    schema,
+    client,
+    databendClient
+  );
+  return db as DatabendDatabase<TSchema, ExtractTablesWithRelations<TSchema>>;
+}
+
+async function createFromDsn<
+  TSchema extends Record<string, unknown> = Record<string, never>,
+>(
+  dsn: string,
+  config: DatabendDrizzleConfig<TSchema> = {}
+): Promise<DatabendDatabase<TSchema, ExtractTablesWithRelations<TSchema>>> {
+  const databendClient = new Client(dsn);
+
+  if (config.pool === false) {
+    const connection = await databendClient.getConn();
+    return createFromClient(connection, config, databendClient);
+  }
+
+  const poolSize = config.pool?.size ?? 4;
+  const pool = createDatabendConnectionPool(databendClient, { size: poolSize });
+  return createFromClient(pool, config, databendClient);
+}
+
+// Overload 1: DSN string (async, auto-pools)
+export function drizzle<
+  TSchema extends Record<string, unknown> = Record<string, never>,
+>(
+  dsn: string
+): Promise<DatabendDatabase<TSchema, ExtractTablesWithRelations<TSchema>>>;
+
+// Overload 2: DSN string + config (async, auto-pools)
+export function drizzle<
+  TSchema extends Record<string, unknown> = Record<string, never>,
+>(
+  dsn: string,
+  config: DatabendDrizzleConfig<TSchema>
+): Promise<DatabendDatabase<TSchema, ExtractTablesWithRelations<TSchema>>>;
+
+// Overload 3: Config with connection (async, auto-pools)
+export function drizzle<
+  TSchema extends Record<string, unknown> = Record<string, never>,
+>(
+  config: DatabendDrizzleConfigWithConnection<TSchema>
+): Promise<DatabendDatabase<TSchema, ExtractTablesWithRelations<TSchema>>>;
+
+// Overload 4: Config with explicit client (sync)
+export function drizzle<
+  TSchema extends Record<string, unknown> = Record<string, never>,
+>(
+  config: DatabendDrizzleConfigWithClient<TSchema>
+): DatabendDatabase<TSchema, ExtractTablesWithRelations<TSchema>>;
+
+// Overload 5: Explicit client (sync, backward compatible)
+export function drizzle<
+  TSchema extends Record<string, unknown> = Record<string, never>,
+>(
+  client: DatabendClientLike,
+  config?: DatabendDrizzleConfig<TSchema>
+): DatabendDatabase<TSchema, ExtractTablesWithRelations<TSchema>>;
+
+// Implementation
+export function drizzle<
+  TSchema extends Record<string, unknown> = Record<string, never>,
+>(
+  clientOrConfigOrDsn:
+    | string
+    | DatabendClientLike
+    | DatabendDrizzleConfigWithConnection<TSchema>
+    | DatabendDrizzleConfigWithClient<TSchema>,
+  config?: DatabendDrizzleConfig<TSchema>
+):
+  | DatabendDatabase<TSchema, ExtractTablesWithRelations<TSchema>>
+  | Promise<DatabendDatabase<TSchema, ExtractTablesWithRelations<TSchema>>> {
+  // String DSN -> async with auto-pool
+  if (typeof clientOrConfigOrDsn === 'string') {
+    return createFromDsn(clientOrConfigOrDsn, config);
+  }
+
+  // Config object with connection or client
+  if (isConfigObject(clientOrConfigOrDsn)) {
+    const configObj = clientOrConfigOrDsn as
+      | DatabendDrizzleConfigWithConnection<TSchema>
+      | DatabendDrizzleConfigWithClient<TSchema>;
+
+    if ('connection' in configObj) {
+      const connConfig =
+        configObj as DatabendDrizzleConfigWithConnection<TSchema>;
+      const { connection, ...restConfig } = connConfig;
+      return createFromDsn(
+        connection,
+        restConfig as DatabendDrizzleConfig<TSchema>
+      );
+    }
+
+    if ('client' in configObj) {
+      const clientConfig = configObj as DatabendDrizzleConfigWithClient<TSchema>;
+      const { client: clientValue, ...restConfig } = clientConfig;
+      return createFromClient(
+        clientValue,
+        restConfig as DatabendDrizzleConfig<TSchema>
+      );
+    }
+
+    throw new Error(
+      'Invalid drizzle config: either connection or client must be provided'
+    );
+  }
+
+  // Direct client (backward compatible)
+  return createFromClient(clientOrConfigOrDsn as DatabendClientLike, config);
+}
+
+export class DatabendDatabase<
+  TFullSchema extends Record<string, unknown> = Record<string, never>,
+  TSchema extends TablesRelationalConfig =
+    ExtractTablesWithRelations<TFullSchema>,
+> extends PgDatabase<DatabendQueryResultHKT, TFullSchema, TSchema> {
+  static readonly [entityKind]: string = 'DatabendDatabase';
+
+  /** The underlying connection or pool */
+  readonly $client: DatabendClientLike;
+
+  /** The Databend Client instance (when created from DSN) */
+  readonly $databendClient?: Client;
+
+  constructor(
+    readonly dialect: DatabendDialect,
+    readonly session: DatabendSession<TFullSchema, TSchema>,
+    schema: RelationalSchemaConfig<TSchema> | undefined,
+    client: DatabendClientLike,
+    databendClient?: Client
+  ) {
+    super(dialect, session, schema);
+    this.$client = client;
+    this.$databendClient = databendClient;
+  }
+
+  async close(): Promise<void> {
+    if (isPool(this.$client) && this.$client.close) {
+      await this.$client.close();
+    }
+    if (!isPool(this.$client)) {
+      await closeClientConnection(this.$client as Connection);
+    }
+  }
+
+  override async transaction<T>(
+    transaction: (tx: DatabendTransaction<TFullSchema, TSchema>) => Promise<T>
+  ): Promise<T> {
+    return await this.session.transaction<T>(transaction);
+  }
+}

package/src/index.ts

@@ -0,0 +1,6 @@
+export * from './driver.ts';
+export * from './session.ts';
+export * from './columns.ts';
+export * from './migrator.ts';
+export * from './client.ts';
+export * from './pool.ts';

package/src/migrator.ts

@@ -0,0 +1,22 @@
+import type { MigrationConfig } from 'drizzle-orm/migrator';
+import { readMigrationFiles } from 'drizzle-orm/migrator';
+import type { DatabendDatabase } from './driver.ts';
+import type { PgSession } from 'drizzle-orm/pg-core/session';
+
+export type DatabendMigrationConfig = MigrationConfig | string;
+
+export async function migrate<TSchema extends Record<string, unknown>>(
+  db: DatabendDatabase<TSchema>,
+  config: DatabendMigrationConfig
+) {
+  const migrationConfig: MigrationConfig =
+    typeof config === 'string' ? { migrationsFolder: config } : config;
+
+  const migrations = readMigrationFiles(migrationConfig);
+
+  await db.dialect.migrate(
+    migrations,
+    db.session as unknown as PgSession,
+    migrationConfig
+  );
+}

package/src/pool.ts

@@ -0,0 +1,233 @@
+import type { Client, Connection } from 'databend-driver';
+import { closeClientConnection, type DatabendConnectionPool } from './client.ts';
+
+export interface DatabendPoolConfig {
+  /** Maximum concurrent connections. Defaults to 4. */
+  size?: number;
+}
+
+export interface DatabendConnectionPoolOptions {
+  /** Maximum concurrent connections. Defaults to 4. */
+  size?: number;
+  /** Timeout in milliseconds to wait for a connection. Defaults to 30000 (30s). */
+  acquireTimeout?: number;
+  /** Maximum number of requests waiting for a connection. Defaults to 100. */
+  maxWaitingRequests?: number;
+  /** Max time (ms) a connection may live before being recycled. */
+  maxLifetimeMs?: number;
+  /** Max idle time (ms) before an idle connection is discarded. */
+  idleTimeoutMs?: number;
+}
+
+export function createDatabendConnectionPool(
+  client: Client,
+  options: DatabendConnectionPoolOptions = {}
+): DatabendConnectionPool & { size: number } {
+  const size = options.size && options.size > 0 ? options.size : 4;
+  const acquireTimeout = options.acquireTimeout ?? 30_000;
+  const maxWaitingRequests = options.maxWaitingRequests ?? 100;
+  const maxLifetimeMs = options.maxLifetimeMs;
+  const idleTimeoutMs = options.idleTimeoutMs;
+  const metadata = new WeakMap<
+    Connection,
+    { createdAt: number; lastUsedAt: number }
+  >();
+
+  type PooledConnection = {
+    connection: Connection;
+    createdAt: number;
+    lastUsedAt: number;
+  };
+
+  const idle: PooledConnection[] = [];
+  const waiting: Array<{
+    resolve: (conn: Connection) => void;
+    reject: (error: Error) => void;
+    timeoutId: ReturnType<typeof setTimeout>;
+  }> = [];
+  let total = 0;
+  let closed = false;
+  let pendingAcquires = 0;
+
+  const shouldRecycle = (conn: PooledConnection, now: number): boolean => {
+    if (maxLifetimeMs !== undefined && now - conn.createdAt >= maxLifetimeMs) {
+      return true;
+    }
+    if (idleTimeoutMs !== undefined && now - conn.lastUsedAt >= idleTimeoutMs) {
+      return true;
+    }
+    return false;
+  };
+
+  const acquire = async (): Promise<Connection> => {
+    if (closed) {
+      throw new Error('Databend connection pool is closed');
+    }
+
+    while (idle.length > 0) {
+      const pooled = idle.pop() as PooledConnection;
+      const now = Date.now();
+      if (shouldRecycle(pooled, now)) {
+        await closeClientConnection(pooled.connection);
+        total = Math.max(0, total - 1);
+        metadata.delete(pooled.connection);
+        continue;
+      }
+      pooled.lastUsedAt = now;
+      metadata.set(pooled.connection, {
+        createdAt: pooled.createdAt,
+        lastUsedAt: pooled.lastUsedAt,
+      });
+      return pooled.connection;
+    }
+
+    if (total < size) {
+      pendingAcquires += 1;
+      total += 1;
+      try {
+        const connection = await client.getConn();
+        if (closed) {
+          await closeClientConnection(connection);
+          total -= 1;
+          throw new Error('Databend connection pool is closed');
+        }
+        const now = Date.now();
+        metadata.set(connection, { createdAt: now, lastUsedAt: now });
+        return connection;
+      } catch (error) {
+        total -= 1;
+        throw error;
+      } finally {
+        pendingAcquires -= 1;
+      }
+    }
+
+    if (waiting.length >= maxWaitingRequests) {
+      throw new Error(
+        `Databend connection pool queue is full (max ${maxWaitingRequests} waiting requests)`
+      );
+    }
+
+    return await new Promise((resolve, reject) => {
+      const timeoutId = setTimeout(() => {
+        const idx = waiting.findIndex((w) => w.timeoutId === timeoutId);
+        if (idx !== -1) {
+          waiting.splice(idx, 1);
+        }
+        reject(
+          new Error(
+            `Databend connection pool acquire timeout after ${acquireTimeout}ms`
+          )
+        );
+      }, acquireTimeout);
+
+      waiting.push({ resolve, reject, timeoutId });
+    });
+  };
+
+  const release = async (connection: Connection): Promise<void> => {
+    const waiter = waiting.shift();
+    if (waiter) {
+      clearTimeout(waiter.timeoutId);
+      const now = Date.now();
+      const meta =
+        metadata.get(connection) ??
+        ({ createdAt: now, lastUsedAt: now } as {
+          createdAt: number;
+          lastUsedAt: number;
+        });
+
+      const expired =
+        maxLifetimeMs !== undefined && now - meta.createdAt >= maxLifetimeMs;
+
+      if (closed) {
+        await closeClientConnection(connection);
+        total = Math.max(0, total - 1);
+        metadata.delete(connection);
+        waiter.reject(new Error('Databend connection pool is closed'));
+        return;
+      }
+
+      if (expired) {
+        await closeClientConnection(connection);
+        total = Math.max(0, total - 1);
+        metadata.delete(connection);
+        try {
+          const replacement = await acquire();
+          waiter.resolve(replacement);
+        } catch (error) {
+          waiter.reject(error as Error);
+        }
+        return;
+      }
+
+      meta.lastUsedAt = now;
+      metadata.set(connection, meta);
+      waiter.resolve(connection);
+      return;
+    }
+
+    if (closed) {
+      await closeClientConnection(connection);
+      metadata.delete(connection);
+      total = Math.max(0, total - 1);
+      return;
+    }
+
+    const now = Date.now();
+    const existingMeta =
+      metadata.get(connection) ??
+      ({ createdAt: now, lastUsedAt: now } as {
+        createdAt: number;
+        lastUsedAt: number;
+      });
+    existingMeta.lastUsedAt = now;
+    metadata.set(connection, existingMeta);
+
+    if (
+      maxLifetimeMs !== undefined &&
+      now - existingMeta.createdAt >= maxLifetimeMs
+    ) {
+      await closeClientConnection(connection);
+      total -= 1;
+      metadata.delete(connection);
+      return;
+    }
+
+    idle.push({
+      connection,
+      createdAt: existingMeta.createdAt,
+      lastUsedAt: existingMeta.lastUsedAt,
+    });
+  };
+
+  const close = async (): Promise<void> => {
+    closed = true;
+
+    const waiters = waiting.splice(0, waiting.length);
+    for (const waiter of waiters) {
+      clearTimeout(waiter.timeoutId);
+      waiter.reject(new Error('Databend connection pool is closed'));
+    }
+
+    const toClose = idle.splice(0, idle.length);
+    await Promise.allSettled(
+      toClose.map((item) => closeClientConnection(item.connection))
+    );
+    total = Math.max(0, total - toClose.length);
+    toClose.forEach((item) => metadata.delete(item.connection));
+
+    const maxWait = 5000;
+    const start = Date.now();
+    while (pendingAcquires > 0 && Date.now() - start < maxWait) {
+      await new Promise((r) => setTimeout(r, 10));
+    }
+  };
+
+  return {
+    acquire,
+    release,
+    close,
+    size,
+  };
+}