import { sqlite } from "https://esm.town/v/std/sqlite";
import { type ResultSet } from "npm:@libsql/client";
import { DataType, Float64, Int32, Table, tableToIPC, Utf8, vectorFromArray } from "npm:apache-arrow";
import { zip } from "npm:lodash-es@4.17.21";
function sqliteTypeToArrowType(sqliteType: string): DataType {
switch (sqliteType.toUpperCase()) {
case "INTEGER":
case "NUMBER":
return new Int32();
case "TEXT":
return new Utf8();
case "REAL":
return new Float64();
default:
return new Utf8();
}
}
export type SQLiteTableExportFormat = "json" | "arrowIPC";
export function resultSetToJSON(res: ResultSet) {
return res.rows.map(row => Object.fromEntries(zip(res.columns, row)));
}
export function resultSetToArrowTable(res: ResultSet) {
const tableData = {};
res.columns.forEach((column, i) => {
const data = res.rows.map(row => row[i]);
tableData[column] = vectorFromArray(data, sqliteTypeToArrowType(res.columnTypes[i]));
});
return new Table(tableData);
}
export async function exportSQLiteTable(
tableName: string,
format: SQLiteTableExportFormat,
limit: number | null,
offset: number | null,
): Promise<Record<string, unknown>[] | Uint8Array> {
const data = await sqlite.execute(`
SELECT * FROM ${tableName} LIMIT ${limit || -1} OFFSET ${offset || 0}
`) as ResultSet;
const rows = data.rows;
const schemaInfo = await sqlite.execute(`PRAGMA table_info(${tableName})`) as ResultSet;
const columns = schemaInfo.rows.map(col => ({ name: String(col[1]), type: sqliteTypeToArrowType(String(col[2])) }));
if (format == "json") {
return resultSetToJSON(data);
}
const arrowTable = resultSetToArrowTable(data);
return tableToIPC(arrowTable);
}