connection: rewrote query execution, batching, added generic sqlite, added node/bun sqlite, aligned repo/mutator results

app/src/data/connection/Connection.ts

@@ -2,12 +2,15 @@ import {
    type AliasableExpression,
    type ColumnBuilderCallback,
    type ColumnDataType,
+   type Compilable,
+   type CompiledQuery,
    type DatabaseIntrospector,
    type Dialect,
    type Expression,
    type Kysely,
    type KyselyPlugin,
    type OnModifyForeignAction,
+   type QueryResult,
    type RawBuilder,
    type SelectQueryBuilder,
    type SelectQueryNode,
@@ -15,7 +18,8 @@ import {
    sql,
 } from "kysely";
 import type { BaseIntrospector, BaseIntrospectorConfig } from "./BaseIntrospector";
-import type { Constructor } from "core";
+import type { Constructor, DB } from "core";
+import { KyselyPluginRunner } from "data/plugins/KyselyPluginRunner";
 
 export type QB = SelectQueryBuilder<any, any, any>;
 
@@ -75,22 +79,44 @@ export type DbFunctions = {
    >;
 };
 
+export type ConnQuery = CompiledQuery | Compilable;
+
+export type ConnQueryResult<T extends ConnQuery> = T extends CompiledQuery<infer R>
+   ? QueryResult<R>
+   : T extends Compilable<infer R>
+     ? QueryResult<R>
+     : never;
+
+export type ConnQueryResults<T extends ConnQuery[]> = {
+   [K in keyof T]: ConnQueryResult<T[K]>;
+};
+
 const CONN_SYMBOL = Symbol.for("bknd:connection");
 
-export abstract class Connection<DB = any> {
+export type Features = {
+   batching: boolean;
+   softscans: boolean;
+};
+
+export abstract class Connection<Client = unknown> {
+   abstract name: string;
    protected initialized = false;
-   kysely: Kysely<DB>;
-   protected readonly supported = {
+   protected pluginRunner: KyselyPluginRunner;
+   protected readonly supported: Partial<Features> = {
       batching: false,
+      softscans: true,
    };
+   kysely: Kysely<DB>;
+   client!: Client;
 
    constructor(
-      kysely: Kysely<DB>,
+      kysely: Kysely<any>,
       public fn: Partial<DbFunctions> = {},
       protected plugins: KyselyPlugin[] = [],
    ) {
       this.kysely = kysely;
       this[CONN_SYMBOL] = true;
+      this.pluginRunner = new KyselyPluginRunner(plugins);
    }
 
    // @todo: consider moving constructor logic here, required by sqlocal
@@ -121,30 +147,46 @@ export abstract class Connection<DB = any> {
       return res.rows.length > 0;
    }
 
-   protected async batch<Queries extends QB[]>(
-      queries: [...Queries],
-   ): Promise<{
-      [K in keyof Queries]: Awaited<ReturnType<Queries[K]["execute"]>>;
-   }> {
-      throw new Error("Batching not supported");
+   protected async transformResultRows(result: any[]): Promise<any[]> {
+      return await this.pluginRunner.transformResultRows(result);
    }
 
-   async batchQuery<Queries extends QB[]>(
-      queries: [...Queries],
-   ): Promise<{
-      [K in keyof Queries]: Awaited<ReturnType<Queries[K]["execute"]>>;
-   }> {
-      // bypass if no client support
-      if (!this.supports("batching")) {
-         const data: any = [];
-         for (const q of queries) {
-            const result = await q.execute();
-            data.push(result);
-         }
-         return data;
-      }
+   /**
+    * Execute a query and return the result including all metadata
+    * returned from the dialect.
+    */
+   async executeQueries<O extends ConnQuery[]>(...qbs: O): Promise<ConnQueryResults<O>> {
+      return Promise.all(qbs.map(async (qb) => await this.kysely.executeQuery(qb))) as any;
+   }
 
-      return await this.batch(queries);
+   async executeQuery<O extends ConnQuery>(qb: O): Promise<ConnQueryResult<O>> {
+      const res = await this.executeQueries(qb);
+      return res[0] as any;
+   }
+
+   protected getCompiled(...qbs: ConnQuery[]): CompiledQuery[] {
+      return qbs.map((qb) => {
+         if ("compile" in qb) {
+            return qb.compile();
+         }
+         return qb;
+      });
+   }
+
+   protected async withTransformedRows<
+      Key extends string = "rows",
+      O extends { [K in Key]: any[] }[] = [],
+   >(result: O, _key?: Key): Promise<O> {
+      return (await Promise.all(
+         result.map(async (row) => {
+            const key = _key ?? "rows";
+            const { [key]: rows, ...r } = row;
+            return {
+               ...r,
+               rows: await this.transformResultRows(rows),
+            };
+         }),
+      )) as any;
    }
 
    protected validateFieldSpecType(type: string): type is FieldSpec["type"] {

app/src/data/connection/connection-test-suite.ts

@@ -0,0 +1,187 @@
+import type { TestRunner } from "core/test";
+import { Connection, type FieldSpec } from "./Connection";
+
+export function connectionTestSuite(
+   testRunner: TestRunner,
+   {
+      makeConnection,
+      rawDialectDetails,
+   }: {
+      makeConnection: () => Connection;
+      rawDialectDetails: string[];
+   },
+) {
+   const { test, expect, describe } = testRunner;
+
+   test("pings", async () => {
+      const connection = makeConnection();
+      const res = await connection.ping();
+      expect(res).toBe(true);
+   });
+
+   test("initializes", async () => {
+      const connection = makeConnection();
+      await connection.init();
+      // @ts-expect-error
+      expect(connection.initialized).toBe(true);
+      expect(connection.client).toBeDefined();
+   });
+
+   test("isConnection", async () => {
+      const connection = makeConnection();
+      expect(Connection.isConnection(connection)).toBe(true);
+   });
+
+   test("getFieldSchema", async () => {
+      const c = makeConnection();
+      const specToNode = (spec: FieldSpec) => {
+         // @ts-expect-error
+         const schema = c.kysely.schema.createTable("test").addColumn(...c.getFieldSchema(spec));
+         return schema.toOperationNode();
+      };
+
+      {
+         // primary
+         const node = specToNode({
+            type: "integer",
+            name: "id",
+            primary: true,
+         });
+         const col = node.columns[0]!;
+         expect(col.primaryKey).toBe(true);
+         expect(col.notNull).toBe(true);
+      }
+
+      {
+         // normal
+         const node = specToNode({
+            type: "text",
+            name: "text",
+         });
+         const col = node.columns[0]!;
+         expect(!col.primaryKey).toBe(true);
+         expect(!col.notNull).toBe(true);
+      }
+
+      {
+         // nullable (expect to be same as normal)
+         const node = specToNode({
+            type: "text",
+            name: "text",
+            nullable: true,
+         });
+         const col = node.columns[0]!;
+         expect(!col.primaryKey).toBe(true);
+         expect(!col.notNull).toBe(true);
+      }
+   });
+
+   describe("schema", async () => {
+      const connection = makeConnection();
+      const fields = [
+         {
+            type: "integer",
+            name: "id",
+            primary: true,
+         },
+         {
+            type: "text",
+            name: "text",
+         },
+         {
+            type: "json",
+            name: "json",
+         },
+      ] as const satisfies FieldSpec[];
+
+      let b = connection.kysely.schema.createTable("test");
+      for (const field of fields) {
+         // @ts-expect-error
+         b = b.addColumn(...connection.getFieldSchema(field));
+      }
+      await b.execute();
+
+      // add index
+      await connection.kysely.schema.createIndex("test_index").on("test").columns(["id"]).execute();
+
+      test("executes query", async () => {
+         await connection.kysely
+            .insertInto("test")
+            .values({ id: 1, text: "test", json: JSON.stringify({ a: 1 }) })
+            .execute();
+
+         const expected = { id: 1, text: "test", json: { a: 1 } };
+
+         const qb = connection.kysely.selectFrom("test").selectAll();
+         const res = await connection.executeQuery(qb);
+         expect(res.rows).toEqual([expected]);
+         expect(rawDialectDetails.every((detail) => detail in res)).toBe(true);
+
+         {
+            const res = await connection.executeQueries(qb, qb);
+            expect(res.length).toBe(2);
+            res.map((r) => {
+               expect(r.rows).toEqual([expected]);
+               expect(rawDialectDetails.every((detail) => detail in r)).toBe(true);
+            });
+         }
+      });
+
+      test("introspects", async () => {
+         const tables = await connection.getIntrospector().getTables({
+            withInternalKyselyTables: false,
+         });
+         const clean = tables.map((t) => ({
+            ...t,
+            columns: t.columns.map((c) => ({
+               ...c,
+               dataType: undefined,
+            })),
+         }));
+
+         expect(clean).toEqual([
+            {
+               name: "test",
+               isView: false,
+               columns: [
+                  {
+                     name: "id",
+                     dataType: undefined,
+                     isNullable: false,
+                     isAutoIncrementing: true,
+                     hasDefaultValue: false,
+                  },
+                  {
+                     name: "text",
+                     dataType: undefined,
+                     isNullable: true,
+                     isAutoIncrementing: false,
+                     hasDefaultValue: false,
+                  },
+                  {
+                     name: "json",
+                     dataType: undefined,
+                     isNullable: true,
+                     isAutoIncrementing: false,
+                     hasDefaultValue: false,
+                  },
+               ],
+            },
+         ]);
+      });
+
+      expect(await connection.getIntrospector().getIndices()).toEqual([
+         {
+            name: "test_index",
+            table: "test",
+            isUnique: false,
+            columns: [
+               {
+                  name: "id",
+                  order: 0,
+               },
+            ],
+         },
+      ]);
+   });
+}

app/src/data/connection/sqlite/GenericSqliteConnection.ts

@@ -0,0 +1,37 @@
+import type { KyselyPlugin } from "kysely";
+import {
+   type IGenericSqlite,
+   type OnCreateConnection,
+   type Promisable,
+   parseBigInt,
+   buildQueryFn,
+   GenericSqliteDialect,
+} from "kysely-generic-sqlite";
+import { SqliteConnection } from "./SqliteConnection";
+
+export type GenericSqliteConnectionConfig = {
+   name: string;
+   additionalPlugins?: KyselyPlugin[];
+   excludeTables?: string[];
+   onCreateConnection?: OnCreateConnection;
+};
+
+export { parseBigInt, buildQueryFn, GenericSqliteDialect, type IGenericSqlite };
+
+export class GenericSqliteConnection<DB = unknown> extends SqliteConnection<DB> {
+   override name = "generic-sqlite";
+
+   constructor(
+      db: DB,
+      executor: () => Promisable<IGenericSqlite>,
+      config?: GenericSqliteConnectionConfig,
+   ) {
+      super({
+         dialect: GenericSqliteDialect,
+         dialectArgs: [executor, config?.onCreateConnection],
+         additionalPlugins: config?.additionalPlugins,
+         excludeTables: config?.excludeTables,
+      });
+      this.client = db;
+   }
+}

app/src/data/connection/sqlite/LibsqlConnection.spec.ts

@@ -0,0 +1,11 @@
+import { connectionTestSuite } from "../connection-test-suite";
+import { LibsqlConnection } from "./LibsqlConnection";
+import { bunTestRunner } from "adapter/bun/test";
+import { describe } from "bun:test";
+
+describe("LibsqlConnection", () => {
+   connectionTestSuite(bunTestRunner, {
+      makeConnection: () => new LibsqlConnection({ url: ":memory:" }),
+      rawDialectDetails: ["rowsAffected", "lastInsertRowid"],
+   });
+});

app/src/data/connection/sqlite/LibsqlConnection.ts

@@ -1,40 +1,26 @@
-import { type Client, type Config, type InStatement, createClient } from "@libsql/client";
+import { createClient, type Client, type Config, type InStatement } from "@libsql/client";
 import { LibsqlDialect } from "@libsql/kysely-libsql";
-import { FilterNumericKeysPlugin } from "data/plugins/FilterNumericKeysPlugin";
-import { KyselyPluginRunner } from "data/plugins/KyselyPluginRunner";
-import { type DatabaseIntrospector, Kysely, ParseJSONResultsPlugin } from "kysely";
-import type { QB } from "../Connection";
-import { SqliteConnection } from "./SqliteConnection";
-import { SqliteIntrospector } from "./SqliteIntrospector";
 import { $console } from "core";
+import { FilterNumericKeysPlugin } from "data/plugins/FilterNumericKeysPlugin";
+import type { ConnQuery, ConnQueryResults } from "../Connection";
+import { SqliteConnection } from "./SqliteConnection";
 
 export const LIBSQL_PROTOCOLS = ["wss", "https", "libsql"] as const;
 export type LibSqlCredentials = Config & {
    protocol?: (typeof LIBSQL_PROTOCOLS)[number];
 };
 
-const plugins = [new FilterNumericKeysPlugin(), new ParseJSONResultsPlugin()];
-
-class CustomLibsqlDialect extends LibsqlDialect {
-   override createIntrospector(db: Kysely<any>): DatabaseIntrospector {
-      return new SqliteIntrospector(db, {
-         excludeTables: ["libsql_wasm_func_table"],
-         plugins,
-      });
-   }
-}
-
-export class LibsqlConnection extends SqliteConnection {
-   private client: Client;
+export class LibsqlConnection extends SqliteConnection<Client> {
+   override name = "libsql";
    protected override readonly supported = {
       batching: true,
+      softscans: true,
    };
 
    constructor(client: Client);
    constructor(credentials: LibSqlCredentials);
    constructor(clientOrCredentials: Client | LibSqlCredentials) {
       let client: Client;
-      let batching_enabled = true;
       if (clientOrCredentials && "url" in clientOrCredentials) {
          let { url, authToken, protocol } = clientOrCredentials;
          if (protocol && LIBSQL_PROTOCOLS.includes(protocol)) {
@@ -48,45 +34,25 @@ export class LibsqlConnection extends SqliteConnection {
          client = clientOrCredentials;
       }
 
-      const kysely = new Kysely({
-         // @ts-expect-error libsql has type issues
-         dialect: new CustomLibsqlDialect({ client }),
-         plugins,
+      super({
+         excludeTables: ["libsql_wasm_func_table"],
+         dialect: LibsqlDialect,
+         dialectArgs: [{ client }],
+         additionalPlugins: [new FilterNumericKeysPlugin()],
       });
 
-      super(kysely, {}, plugins);
       this.client = client;
-      this.supported.batching = batching_enabled;
    }
 
-   getClient(): Client {
-      return this.client;
-   }
-
-   protected override async batch<Queries extends QB[]>(
-      queries: [...Queries],
-   ): Promise<{
-      [K in keyof Queries]: Awaited<ReturnType<Queries[K]["execute"]>>;
-   }> {
-      const stms: InStatement[] = queries.map((q) => {
-         const compiled = q.compile();
+   override async executeQueries<O extends ConnQuery[]>(...qbs: O): Promise<ConnQueryResults<O>> {
+      const compiled = this.getCompiled(...qbs);
+      const stms: InStatement[] = compiled.map((q) => {
          return {
-            sql: compiled.sql,
-            args: compiled.parameters as any[],
+            sql: q.sql,
+            args: q.parameters as any[],
          };
       });
 
-      const res = await this.client.batch(stms);
-
-      // let it run through plugins
-      const kyselyPlugins = new KyselyPluginRunner(this.plugins);
-
-      const data: any = [];
-      for (const r of res) {
-         const rows = await kyselyPlugins.transformResultRows(r.rows);
-         data.push(rows);
-      }
-
-      return data;
+      return this.withTransformedRows(await this.client.batch(stms)) as any;
    }
 }

app/src/data/connection/sqlite/SqliteConnection.ts

@@ -1,16 +1,49 @@
-import type { ColumnDataType, ColumnDefinitionBuilder, Kysely, KyselyPlugin } from "kysely";
+import {
+   ParseJSONResultsPlugin,
+   type ColumnDataType,
+   type ColumnDefinitionBuilder,
+   type Dialect,
+   Kysely,
+   type KyselyPlugin,
+} from "kysely";
 import { jsonArrayFrom, jsonBuildObject, jsonObjectFrom } from "kysely/helpers/sqlite";
 import { Connection, type DbFunctions, type FieldSpec, type SchemaResponse } from "../Connection";
+import type { Constructor } from "core";
+import { customIntrospector } from "../Connection";
+import { SqliteIntrospector } from "./SqliteIntrospector";
+
+export type SqliteConnectionConfig<
+   CustomDialect extends Constructor<Dialect> = Constructor<Dialect>,
+> = {
+   excludeTables?: string[];
+   dialect: CustomDialect;
+   dialectArgs?: ConstructorParameters<CustomDialect>;
+   additionalPlugins?: KyselyPlugin[];
+   customFn?: Partial<DbFunctions>;
+};
+
+export abstract class SqliteConnection<Client = unknown> extends Connection<Client> {
+   override name = "sqlite";
+
+   constructor(config: SqliteConnectionConfig) {
+      const { excludeTables, dialect, dialectArgs = [], additionalPlugins } = config;
+      const plugins = [new ParseJSONResultsPlugin(), ...(additionalPlugins ?? [])];
+
+      const kysely = new Kysely({
+         dialect: customIntrospector(dialect, SqliteIntrospector, {
+            excludeTables,
+            plugins,
+         }).create(...dialectArgs),
+         plugins,
+      });
 
-export class SqliteConnection extends Connection {
-   constructor(kysely: Kysely<any>, fn: Partial<DbFunctions> = {}, plugins: KyselyPlugin[] = []) {
       super(
          kysely,
          {
-            ...fn,
             jsonArrayFrom,
             jsonObjectFrom,
             jsonBuildObject,
+            ...(config.customFn ?? {}),
          },
          plugins,
       );
@@ -43,7 +76,7 @@ export class SqliteConnection extends Connection {
                if (spec.onUpdate) relCol = relCol.onUpdate(spec.onUpdate);
                return relCol;
             }
-            return spec.nullable ? col : col.notNull();
+            return col;
          },
       ] as const;
    }

app/src/data/connection/sqlite/SqliteLocalConnection.ts

@@ -1,31 +1,14 @@
-import {
-   type DatabaseIntrospector,
-   Kysely,
-   ParseJSONResultsPlugin,
-   type SqliteDatabase,
-   SqliteDialect,
-} from "kysely";
+import { type SqliteDatabase, SqliteDialect } from "kysely";
 import { SqliteConnection } from "./SqliteConnection";
-import { SqliteIntrospector } from "./SqliteIntrospector";
 
-const plugins = [new ParseJSONResultsPlugin()];
+export class SqliteLocalConnection extends SqliteConnection<SqliteDatabase> {
+   override name = "sqlite-local";
 
-class CustomSqliteDialect extends SqliteDialect {
-   override createIntrospector(db: Kysely<any>): DatabaseIntrospector {
-      return new SqliteIntrospector(db, {
-         excludeTables: ["test_table"],
-         plugins,
+   constructor(database: SqliteDatabase) {
+      super({
+         dialect: SqliteDialect,
+         dialectArgs: [{ database }],
       });
-   }
-}
-
-export class SqliteLocalConnection extends SqliteConnection {
-   constructor(private database: SqliteDatabase) {
-      const kysely = new Kysely({
-         dialect: new CustomSqliteDialect({ database }),
-         plugins,
-      });
-
-      super(kysely, {}, plugins);
+      this.client = database;
    }
 }