yasumi/packages/data-sources/duckdb/src/index.ts

import { Resource, validation as v, BaseResourceType } from '@modal-sh/yasumi';
import { DataSource, ResourceIdConfig } from '@modal-sh/yasumi/backend';
import { Database } from 'duckdb-async';
import assert from 'assert';

type ID = number;

interface DuckDbDataSourceBase<
  ID,
  Schema extends v.BaseSchema = v.BaseSchema,
  CurrentName extends string = string,
  CurrentRouteName extends string = string,
  Data extends object = v.Output<Schema>,
> extends DataSource<Data, ID> {
  resource?: Resource<BaseResourceType & {
    schema: Schema,
    name: CurrentName,
    routeName: CurrentRouteName,
  }>;

  db?: Database;
}

export const AutoincrementIdConfig = {
  // TODO add options: https://duckdb.org/docs/sql/statements/create_sequence
  generationStrategy: async (dataSourceRaw: DataSource) => {
    const dataSource = dataSourceRaw as DuckDbDataSourceBase<ID>;
    assert(typeof dataSource.db !== 'undefined');
    assert(typeof dataSource.resource !== 'undefined');
    const idAttr = dataSource.resource.state.shared.get('idAttr');
    assert(typeof idAttr === 'string');

    const con = await dataSource.db.connect();
    const stmt = await con.prepare(`
      SELECT nextval('${dataSource.resource.state.routeName}_sequence') as ${idAttr}; 
    `);
    const [v] = await stmt.all();
    return v[idAttr];
  },
  schema: v.number(),
  serialize: (v: unknown) => v?.toString() ?? '',
  deserialize: (v: string) => Number(v),
}

export class DuckDbDataSource<
  Schema extends v.BaseSchema = v.BaseSchema,
  CurrentName extends string = string,
  CurrentRouteName extends string = string,
  Data extends object = v.Output<Schema>,
> implements DuckDbDataSourceBase<ID, Schema, CurrentName, CurrentRouteName, Data> {
  resource?: Resource<BaseResourceType & {
    schema: Schema,
    name: CurrentName,
    routeName: CurrentRouteName,
  }>;

  db?: Database;

  constructor(private readonly path: string) {
    // noop
  }

  async initialize() {
    assert(typeof this.path !== 'undefined');
    assert(typeof this.resource !== 'undefined');

    const idConfig = this.resource.state.shared.get('idConfig') as ResourceIdConfig<any> | undefined;
    assert(typeof idConfig !== 'undefined');

    const idAttr = this.resource.state.shared.get('idAttr');
    assert(typeof idAttr === 'string');

    const idSchema = idConfig.schema as v.BaseSchema;

    this.db = await Database.create(this.path);
    const clause = `CREATE TABLE IF NOT EXISTS ${this.resource.state.routeName}`;
    const resourceSchema = this.resource.schema as unknown as v.ObjectSchema<any>;
    const tableSchema = Object.entries(resourceSchema.entries)
      .map(([columnName, columnDefRaw]) => {
        const columnDef = columnDefRaw as unknown as v.BaseSchema;
        return [columnName, columnDef.type].join(' ');
      })
      .join(',');
    let sequenceSql = '';
    let defaultValue = '';
    let idType = 'STRING';
    if (idSchema.type === 'number') {
      // TODO support more sequence statements: https://duckdb.org/docs/sql/statements/create_sequence
      sequenceSql = `CREATE SEQUENCE IF NOT EXISTS ${this.resource.state.routeName}_sequence START 1;`;
      defaultValue = `DEFAULT nextval('${this.resource.state.routeName}_sequence')`;
      idType = 'INTEGER';
    }
    const sql = `${sequenceSql}${clause} (${idAttr} ${idType} ${defaultValue},${tableSchema});`;
    const con = await this.db.connect();
    const stmt = await con.prepare(sql);
    await stmt.run();
  }

  prepareResource<Schema extends v.BaseSchema>(resource: Resource<BaseResourceType & {
    schema: Schema,
    name: CurrentName,
    routeName: CurrentRouteName,
  }>) {
    resource.dataSource = resource.dataSource ?? this;
    const originalResourceId = resource.id;
    resource.id = <NewIdAttr extends BaseResourceType['idAttr'], NewIdSchema extends BaseResourceType['idSchema']>(newIdAttr: NewIdAttr, params: ResourceIdConfig<NewIdSchema>) => {
      originalResourceId(newIdAttr, params);
      return resource as Resource<BaseResourceType & {
        name: CurrentName,
        routeName: CurrentRouteName,
        schema: Schema,
        idAttr: NewIdAttr,
        idSchema: NewIdSchema,
      }>;
    };
    this.resource = resource as any;
  }

  async getMultiple(query) {
    // TODO translate query to SQL statements
    assert(typeof this.db !== 'undefined');
    assert(typeof this.resource !== 'undefined');

    const con = await this.db.connect();
    const stmt = await con.prepare(`
      SELECT * FROM ${this.resource.state.routeName};
    `);
    const data = await stmt.all();
    return data as Data[];
  }

  async newId() {
    assert(typeof this.resource !== 'undefined');
    const idConfig = this.resource.state.shared.get('idConfig') as ResourceIdConfig<any>;
    assert(typeof idConfig !== 'undefined');

    const theNewId = await idConfig.generationStrategy(this);
    return theNewId as ID;
  }

  async create(data: Data) {
    assert(typeof this.db !== 'undefined');
    assert(typeof this.resource !== 'undefined');

    const idConfig = this.resource.state.shared.get('idConfig') as ResourceIdConfig<any> | undefined;
    assert(typeof idConfig !== 'undefined');

    const idAttr = this.resource.state.shared.get('idAttr');
    assert(typeof idAttr === 'string');

    let theId: any;
    const { [idAttr]: dataId } = data as Record<string, unknown>;
    if (typeof dataId !== 'undefined') {
      theId = idConfig.deserialize((data as Record<string, string>)[idAttr]);
    } else {
      const newId = await this.newId();
      theId = idConfig.deserialize(newId.toString());
    }
    const effectiveData = {
      ...data,
    } as Record<string, unknown>;
    effectiveData[idAttr] = theId;

    const clause = `INSERT INTO ${this.resource.state.routeName}`;
    const keys = Object.keys(effectiveData).join(',');
    const values = Object.values(effectiveData).map((d) => JSON.stringify(d).replace(/"/g, "'")).join(',');
    const sql = `${clause} (${keys}) VALUES (${values});`;
    const con = await this.db.connect();
    const stmt = await con.prepare(sql);
    await stmt.run();
    const newData = {
      ...effectiveData
    };
    return newData as Data;
  }

  async getTotalCount(query) {
    assert(typeof this.db !== 'undefined');
    assert(typeof this.resource !== 'undefined');

    const idAttr = this.resource.state.shared.get('idAttr');
    assert(typeof idAttr === 'string');

    const con = await this.db.connect();
    const stmt = await con.prepare(`
      SELECT COUNT(*) as c FROM ${this.resource.state.routeName};
    `);
    const [data] = await stmt.all();
    return data['c'] as unknown as number;
  }

  async getById(id: ID) {
    assert(typeof this.db !== 'undefined');
    assert(typeof this.resource !== 'undefined');

    const idAttr = this.resource.state.shared.get('idAttr');
    assert(typeof idAttr === 'string');

    const con = await this.db.connect();
    const stmt = await con.prepare(`
      SELECT * FROM ${this.resource.state.routeName} WHERE ${idAttr} = ${JSON.stringify(id).replace(/"/g, "'")};
    `);
    const [data = null] = await stmt.all();
    return data as Data | null;
  }

  async getSingle(query) {
    assert(typeof this.db !== 'undefined');
    assert(typeof this.resource !== 'undefined');

    const con = await this.db.connect();
    const stmt = await con.prepare(`
      SELECT * FROM ${this.resource.state.routeName} LIMIT 1;
    `);
    const [data = null] = await stmt.all();
    return data as Data | null;
  }

  async delete(id: ID) {
    assert(typeof this.db !== 'undefined');
    assert(typeof this.resource !== 'undefined');

    const idAttr = this.resource.state.shared.get('idAttr');
    assert(typeof idAttr === 'string');

    const con = await this.db.connect();
    const stmt = await con.prepare(`
      DELETE FROM ${this.resource.state.routeName} WHERE ${idAttr} = ${JSON.stringify(id).replace(/"/g, "'")};
    `);
    await stmt.run();
  }

  async emplace(id: ID, data: Data) {
    assert(typeof this.db !== 'undefined');
    assert(typeof this.resource !== 'undefined');

    const idAttr = this.resource.state.shared.get('idAttr');
    assert(typeof idAttr === 'string');

    const clause = `INSERT OR REPLACE INTO ${this.resource.state.routeName}`;
    const keys = Object.keys(data).join(',');
    const values = Object.values(data).map((d) => JSON.stringify(d).replace(/"/g, "'")).join(',');
    const sql = `${clause} (${idAttr},${keys}) VALUES (${id},${values});`;
    const con = await this.db.connect();
    const stmt = await con.prepare(sql);
    const [newData] = await stmt.all();
    // TODO check if created flag
    return [newData, false] as [Data, boolean];
  }

  async patch(id: ID, data: Partial<Data>) {
    assert(typeof this.db !== 'undefined');
    assert(typeof this.resource !== 'undefined');

    const idAttr = this.resource.state.shared.get('idAttr');
    assert(typeof idAttr === 'string');

    const clause = `UPDATE ${this.resource.state.routeName}`;
    const setParams = Object.entries(data).map(([key, value]) => (
      `${key} = ${JSON.stringify(value).replace(/"/g, "'")}`
    )).join(',');
    const sql = `${clause} SET ${setParams} WHERE ${idAttr} = ${JSON.stringify(id).replace(/"/g, "'")}`
    const con = await this.db.connect();
    const stmt = await con.prepare(sql);
    const [newData] = await stmt.all();
    return newData as Data;
  }
}