import { KUNKUN_EXT_IDENTIFIER, type IDb, type IKV } from "@kksh/api" import { CmdType, convertDateToSqliteString, Ext, ExtCmd, ExtData, SearchMode, SearchModeEnum, SQLSortOrder, SQLSortOrderEnum } from "@kksh/api/models" import * as relations from "@kksh/drizzle/relations" import * as orm from "drizzle-orm" import type { SelectedFields } from "drizzle-orm/sqlite-core" import * as v from "valibot" import * as schema from "../drizzle/schema" import { db } from "./proxy" /* -------------------------------------------------------------------------- */ /* Built-in Extensions */ /* -------------------------------------------------------------------------- */ export function getExtClipboard(): Promise { return getExtensionByIdentifierExpectExists(KUNKUN_EXT_IDENTIFIER.KUNKUN_CLIPBOARD_EXT_IDENTIFIER) } export function getExtQuickLinks(): Promise { return getExtensionByIdentifierExpectExists( KUNKUN_EXT_IDENTIFIER.KUNKUN_QUICK_LINKS_EXT_IDENTIFIER ) } export function getExtRemote(): Promise { return getExtensionByIdentifierExpectExists(KUNKUN_EXT_IDENTIFIER.KUNKUN_REMOTE_EXT_IDENTIFIER) } export function getExtScriptCmd(): Promise { return getExtensionByIdentifierExpectExists( KUNKUN_EXT_IDENTIFIER.KUNKUN_SCRIPT_CMD_EXT_IDENTIFIER ) } export function getExtDev(): Promise { return getExtensionByIdentifierExpectExists(KUNKUN_EXT_IDENTIFIER.KUNKUN_DEV_EXT_IDENTIFIER) } /* -------------------------------------------------------------------------- */ /* Extension CRUD */ /* -------------------------------------------------------------------------- */ export function createExtension(ext: { identifier: string version: string enabled?: boolean path?: string data?: any }) { return db .insert(schema.extensions) .values({ identifier: ext.identifier, version: ext.version, enabled: ext.enabled, path: ext.path, data: ext.data }) .run() } export async function getUniqueExtensionByIdentifier(identifier: string): Promise { const ext = await db .select() .from(schema.extensions) .where(orm.eq(schema.extensions.identifier, identifier)) .get() const result = v.safeParse(v.optional(Ext), ext) if (!result.success) { console.error("Failed to parse extension:", v.flatten(result.issues)) return undefined } const parsed = result.output return parsed } /** * Use this function when you expect the extension to exist. Such as builtin extensions. * @param identifier * @returns */ export function getExtensionByIdentifierExpectExists(identifier: string): Promise { return getUniqueExtensionByIdentifier(identifier).then((ext) => { if (!ext) { throw new Error(`Unexpexted Error: Extension ${identifier} not found`) } return ext }) } export async function getAllExtensions(): Promise { const exts = await db.select().from(schema.extensions).all() return v.parse(v.array(Ext), exts) } /** * There can be duplicate extensions with the same identifier. Store and Dev extensions can have the same identifier. * But install path must be unique. * @param path */ export async function getUniqueExtensionByPath(path: string) { const ext = await db .select() .from(schema.extensions) .where(orm.eq(schema.extensions.path, path)) .get() return v.parse(Ext, ext) } export function getAllExtensionsByIdentifier(identifier: string): Promise { return db .select() .from(schema.extensions) .where(orm.eq(schema.extensions.identifier, identifier)) .all() .then((exts) => v.parse(v.array(Ext), exts)) } export function deleteExtensionByPath(path: string): Promise { return db .delete(schema.extensions) .where(orm.eq(schema.extensions.path, path)) .run() .then(() => undefined) } export function deleteExtensionByExtId(extId: number): Promise { return db .delete(schema.extensions) .where(orm.eq(schema.extensions.extId, extId)) .run() .then(() => undefined) } /* -------------------------------------------------------------------------- */ /* Extension Command CRUD */ /* -------------------------------------------------------------------------- */ export function createCommand(data: { extId: number name: string cmdType: CmdType data: string alias?: string hotkey?: string enabled?: boolean }) { return db .insert(schema.commands) .values({ extId: data.extId, name: data.name, type: data.cmdType, data: data.data, alias: data.alias, hotkey: data.hotkey, enabled: data.enabled ?? true }) .run() .then(() => undefined) } export async function getCmdById(cmdId: number): Promise { const cmd = await db .select() .from(schema.commands) .where(orm.eq(schema.commands.cmdId, cmdId)) .get() return v.parse(ExtCmd, cmd) } export async function getAllCmds(): Promise { const cmds = await db.select().from(schema.commands).all() return v.parse(v.array(ExtCmd), cmds) } export function getCommandsByExtId(extId: number) { return db .select() .from(schema.commands) .where(orm.eq(schema.commands.extId, extId)) .all() .then((cmds) => v.parse(v.array(ExtCmd), cmds)) } export function deleteCmdById(cmdId: number) { return db .delete(schema.commands) .where(orm.eq(schema.commands.cmdId, cmdId)) .run() .then(() => undefined) } export function updateCmdByID(data: { cmdId: number name: string cmdType: CmdType data: string alias?: string hotkey?: string enabled: boolean }) { return db .update(schema.commands) .set({ name: data.name, type: data.cmdType, data: data.data, alias: data.alias, // optional hotkey: data.hotkey, // optional enabled: data.enabled // in drizzle schema, use integer({ mode: 'boolean' }) for boolean sqlite // enabled: data.enabled ? String(data.enabled) : undefined }) .where(orm.eq(schema.commands.cmdId, data.cmdId)) .run() .then(() => undefined) } /* -------------------------------------------------------------------------- */ /* Extension Data CRUD */ /* -------------------------------------------------------------------------- */ export const ExtDataField = v.union([v.literal("data"), v.literal("search_text")]) export type ExtDataField = v.InferOutput function convertRawExtDataToExtData(rawData?: { createdAt: string updatedAt: string data: null | string searchText?: null | string dataId: number extId: number dataType: string }): ExtData | undefined { if (!rawData) { return rawData } const parsedRes = v.safeParse(ExtData, { ...rawData, createdAt: new Date(rawData.createdAt), updatedAt: new Date(rawData.updatedAt), data: rawData.data ?? undefined, searchText: rawData.searchText ?? undefined }) if (parsedRes.success) { return parsedRes.output } else { console.error("Extension Data Parse Failure", parsedRes.issues) throw new Error("Fail to parse extension data") } } export function createExtensionData(data: { extId: number dataType: string data: string searchText?: string }) { return db.insert(schema.extensionData).values(data).run() } export function getExtensionDataById(dataId: number, fields?: ExtDataField[]) { const _fields = fields ?? [] const selectQuery: SelectedFields = { dataId: schema.extensionData.dataId, extId: schema.extensionData.extId, dataType: schema.extensionData.dataType, metadata: schema.extensionData.metadata, createdAt: schema.extensionData.createdAt, updatedAt: schema.extensionData.updatedAt // data: schema.extensionData.data, // searchText: schema.extensionData.searchText } if (_fields.includes("data")) { selectQuery["data"] = schema.extensionData.data } if (_fields.includes("search_text")) { selectQuery["searchText"] = schema.extensionData.searchText } return db .select(selectQuery) .from(schema.extensionData) .where(orm.eq(schema.extensionData.dataId, dataId)) .get() .then((rawData) => { console.log("Raw Data", rawData) // @ts-expect-error - rawData is unknown, but will be safe parsed with valibot return convertRawExtDataToExtData(rawData) }) } export async function searchExtensionData(searchParams: { extId: number searchMode: SearchMode dataId?: number dataType?: string searchText?: string afterCreatedAt?: string beforeCreatedAt?: string limit?: number offset?: number orderByCreatedAt?: SQLSortOrder orderByUpdatedAt?: SQLSortOrder fields?: ExtDataField[] }): Promise { const fields = v.parse(v.optional(v.array(ExtDataField), []), searchParams.fields) const _fields = fields ?? [] // Build the select query based on fields const selectQuery: SelectedFields = { dataId: schema.extensionData.dataId, extId: schema.extensionData.extId, dataType: schema.extensionData.dataType, createdAt: schema.extensionData.createdAt, updatedAt: schema.extensionData.updatedAt } if (_fields.includes("data")) { selectQuery["data"] = schema.extensionData.data } if (_fields.includes("search_text")) { selectQuery["searchText"] = schema.extensionData.searchText } // Build the query let baseQuery = db.select(selectQuery).from(schema.extensionData) // Add FTS join if needed if (searchParams.searchMode === SearchModeEnum.FTS && searchParams.searchText) { // @ts-expect-error - The join type is correct but TypeScript can't infer it properly baseQuery = baseQuery.innerJoin( schema.extensionDataFts, orm.eq(schema.extensionData.dataId, schema.extensionDataFts.dataId) ) } // Add conditions const conditions = [orm.eq(schema.extensionData.extId, searchParams.extId)] if (searchParams.dataId) { conditions.push(orm.eq(schema.extensionData.dataId, searchParams.dataId)) } if (searchParams.dataType) { conditions.push(orm.eq(schema.extensionData.dataType, searchParams.dataType)) } if (searchParams.searchText) { switch (searchParams.searchMode) { case SearchModeEnum.ExactMatch: conditions.push(orm.eq(schema.extensionData.searchText, searchParams.searchText)) break case SearchModeEnum.Like: conditions.push(orm.like(schema.extensionData.searchText, `%${searchParams.searchText}%`)) break case SearchModeEnum.FTS: conditions.push( orm.sql`${schema.extensionDataFts.searchText} MATCH ${searchParams.searchText}` ) break } } if (searchParams.afterCreatedAt) { conditions.push(orm.gt(schema.extensionData.createdAt, searchParams.afterCreatedAt)) } if (searchParams.beforeCreatedAt) { conditions.push(orm.lt(schema.extensionData.createdAt, searchParams.beforeCreatedAt)) } // Build the final query with all conditions and modifiers const query = baseQuery .where(orm.and(...conditions)) .orderBy( searchParams.orderByCreatedAt ? searchParams.orderByCreatedAt === SQLSortOrderEnum.Asc ? orm.asc(schema.extensionData.createdAt) : orm.desc(schema.extensionData.createdAt) : searchParams.orderByUpdatedAt ? searchParams.orderByUpdatedAt === SQLSortOrderEnum.Asc ? orm.asc(schema.extensionData.updatedAt) : orm.desc(schema.extensionData.updatedAt) : orm.asc(schema.extensionData.createdAt) // Default ordering ) .limit(searchParams.limit ?? 100) // Default limit .offset(searchParams.offset ?? 0) // Default offset // Execute query and convert results const results = await query.all() return results .map((rawData) => { // @ts-expect-error - rawData is unknown, but will be safe parsed with valibot return convertRawExtDataToExtData(rawData) }) .filter((item): item is ExtData => item !== undefined) } export function deleteExtensionDataById(dataId: number) { return db .delete(schema.extensionData) .where(orm.eq(schema.extensionData.dataId, dataId)) .run() .then(() => undefined) } export function updateExtensionDataById(data: { dataId: number data: string searchText?: string }): Promise { return db .update(schema.extensionData) .set({ data: data.data, searchText: data.searchText }) .where(orm.eq(schema.extensionData.dataId, data.dataId)) .run() .then(() => undefined) } /** * Database API for extensions. * Extensions shouldn't have full access to the database, they can only access their own data. * When an extension is loaded, the main thread will create an instance of this class and * expose it to the extension. */ export class JarvisExtDB implements IDb { extId: number constructor(extId: number) { this.extId = extId } async add(data: { data: string; dataType?: string; searchText?: string }): Promise { return createExtensionData({ data: data.data, dataType: data.dataType ?? "default", searchText: data.searchText, extId: this.extId }).then(() => undefined) } async delete(dataId: number): Promise { // Verify if this data belongs to this extension const d = await getExtensionDataById(dataId) if (!d || d.extId !== this.extId) { throw new Error("Extension Data not found") } return await deleteExtensionDataById(dataId) } async search(searchParams: { dataId?: number searchMode?: SearchMode dataType?: string searchText?: string afterCreatedAt?: Date beforeCreatedAt?: Date limit?: number orderByCreatedAt?: SQLSortOrder orderByUpdatedAt?: SQLSortOrder fields?: ExtDataField[] }): Promise { const beforeCreatedAt = searchParams.beforeCreatedAt ? convertDateToSqliteString(searchParams.beforeCreatedAt) : undefined const afterCreatedAt = searchParams.afterCreatedAt ? convertDateToSqliteString(searchParams.afterCreatedAt) : undefined return searchExtensionData({ ...searchParams, searchMode: searchParams.searchMode ?? SearchModeEnum.FTS, extId: this.extId, beforeCreatedAt, afterCreatedAt }) } /** * Retrieve all data of this extension. * Use `search()` method for more advanced search. * @param options optional fields to retrieve. By default, data and searchText are not returned. * @returns */ retrieveAll(options: { fields?: ExtDataField[] }): Promise { return this.search({ fields: options.fields }) } /** * Retrieve all data of this extension by type. * Use `search()` method for more advanced search. * @param dataType * @returns */ retrieveAllByType(dataType: string): Promise { return this.search({ dataType }) } /** * Delete all data of this extension. */ deleteAll(): Promise { return this.search({}) .then((items) => { return Promise.all(items.map((item) => this.delete(item.dataId))) }) .then(() => {}) } /** * Update data and searchText of this extension. * @param dataId unique id of the data * @param data * @param searchText * @returns */ async update(data: { dataId: number; data: string; searchText?: string }): Promise { const d = await getExtensionDataById(data.dataId) if (!d || d.extId !== this.extId) { throw new Error("Extension Data not found") } return updateExtensionDataById(data) } } export class KV implements IKV { extId: number db: JarvisExtDB private DataType: string = "kunkun_kv" constructor(extId: number) { this.extId = extId this.db = new JarvisExtDB(extId) } get(key: string): Promise { return this.db .search({ dataType: this.DataType, searchText: key, searchMode: SearchModeEnum.ExactMatch, fields: ["search_text", "data"] }) .then((items) => { if (items.length === 0) { return null } else if (items.length > 1) { throw new Error("Multiple KVs with the same key") } return items[0]?.data ? (JSON.parse(items[0].data).value as T) : null }) .catch((err) => { console.warn(err) return null }) } set(key: string, value: string): Promise { return this.db .search({ dataType: this.DataType, searchText: key, searchMode: SearchModeEnum.ExactMatch }) .then((items) => { if (items.length === 0) { return this.db.add({ data: JSON.stringify({ value: value }), dataType: this.DataType, searchText: key }) } else if (items.length === 1) { return this.db.update({ dataId: items[0]!.dataId, data: JSON.stringify({ value: value }), searchText: key }) } else { return Promise.all(items.map((item) => this.db.delete(item.dataId))).then(() => Promise.resolve() ) } }) } delete(key: string): Promise { return this.db .search({ dataType: this.DataType, searchText: key, searchMode: SearchModeEnum.ExactMatch }) .then((items) => { return Promise.all(items.map((item) => this.db.delete(item.dataId))).then(() => Promise.resolve() ) }) } exists(key: string): Promise { return this.db .search({ dataType: this.DataType, searchText: key, searchMode: SearchModeEnum.ExactMatch, fields: [] }) .then((items) => { return items.length > 0 }) } }