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, > extends DataSource { resource?: Resource; 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; 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, > implements DuckDbDataSourceBase { resource?: Resource; 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 | 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; 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(resource: Resource) { resource.dataSource = resource.dataSource ?? this; const originalResourceId = resource.id; resource.id = (newIdAttr: NewIdAttr, params: ResourceIdConfig) => { originalResourceId(newIdAttr, params); return resource as Resource; }; 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; 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 | 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; if (typeof dataId !== 'undefined') { theId = idConfig.deserialize((data as Record)[idAttr]); } else { const newId = await this.newId(); theId = idConfig.deserialize(newId.toString()); } const effectiveData = { ...data, } as Record; 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) { 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; } }