language: TypeScript
16 KB / 559 lines / 496 loc
import type { Command } from "@yakatak/operation";
import { deriveUid, type Source } from "@yakatak/shared";
import BetterSqlite3 from "better-sqlite3";
import { Kysely, sql, SqliteDialect } from "kysely";
import fs from "node:fs/promises";
import path from "node:path";
import type { Database as Schema } from "./schema.ts";
export type CollectJob = {
id: number;
domain: string;
request: { type: "page" | "zulip"; url: string };
};
export type PostprocessJob = {
id: number;
detailId: number;
detailImagePath: string;
};
type Card = {
uid: string;
url: string | null;
title: string | null;
numTiles: number;
};
function decodeUid(uid: Buffer): string {
return uid.toString("base64url");
}
function encodeUid(s: string): Buffer {
return Buffer.from(s, "base64url");
}
export async function ensureDatabaseSchema(sqlite: BetterSqlite3.Database): Promise<void> {
const tableExists = sqlite
.prepare<[], { count: number }>(
"SELECT COUNT(*) AS count FROM sqlite_master WHERE type='table' AND name='card'",
)
.get();
if (tableExists!.count > 0) return;
const schemaPath = path.join(import.meta.dirname, "schema.sql");
const schema = await fs.readFile(schemaPath, "utf-8");
sqlite.exec(schema);
console.log("Initialized database schema");
}
export class YakatakDb {
private sqlite: BetterSqlite3.Database;
private db: Kysely<Schema>;
constructor(dbPath: string) {
this.sqlite = new BetterSqlite3(dbPath);
this.sqlite.pragma("foreign_keys = ON");
this.db = new Kysely<Schema>({
dialect: new SqliteDialect({ database: this.sqlite }),
});
}
async init() {
return ensureDatabaseSchema(this.sqlite);
}
async close(): Promise<void> {
await this.db.destroy();
}
async enqueueCard(url: string): Promise<number> {
const { uid, serialized: source } = await deriveUid({ type: "page", url });
const domain = new URL(url).hostname;
return this.db.transaction().execute(async (trx) => {
const existingCard = await trx
.selectFrom("card")
.select("id")
.where("source", "=", sql<Buffer>`jsonb(${source})`)
.executeTakeFirst();
if (existingCard) return existingCard.id;
const card = await trx
.insertInto("card")
.values({
uid: Buffer.from(uid),
source: sql<Buffer>`jsonb(${source})`,
url,
})
.returning("id")
.executeTakeFirstOrThrow();
await trx
.insertInto("collect_job")
.values({
request: sql<Buffer>`jsonb(${source})`,
domain,
})
.execute();
return card.id;
});
}
async expireDomainTokens(): Promise<void> {
await this.db
.deleteFrom("domain_token_lease")
.where("leased_until", "<=", sql<string>`datetime('now')`)
.execute();
}
async existsUnclaimedCollectJob(): Promise<boolean> {
const row = await this.db
.selectFrom("collect_job")
.select(sql<number>`1`.as("one"))
.where("claimed_at", "is", null)
.limit(1)
.executeTakeFirst();
return row != null;
}
async claimCollectJob(
claimedBy: string,
tokensPerDomain: number,
leaseDurationSec: number,
): Promise<CollectJob | undefined> {
return this.db.transaction().execute(async (trx) => {
const result = await sql<{ id: number; request: string; domain: string }>`
WITH used_tokens AS (
SELECT domain, COUNT(*) as count
FROM domain_token_lease
GROUP BY domain
)
UPDATE collect_job
SET
claimed_by = ${claimedBy},
claimed_at = datetime('now')
WHERE id = (
SELECT candidate.id
FROM collect_job candidate
LEFT JOIN used_tokens ON used_tokens.domain = candidate.domain
WHERE candidate.claimed_at IS NULL
AND COALESCE(used_tokens.count, 0) < ${tokensPerDomain}
ORDER BY candidate.created_at ASC
LIMIT 1
)
RETURNING id, json(request) AS request, domain
`.execute(trx);
const row = result.rows[0];
if (!row) return undefined;
await trx
.insertInto("domain_token_lease")
.values({
domain: row.domain,
leased_until: sql<string>`datetime('now', '+' || ${leaseDurationSec} || ' seconds')`,
})
.execute();
return { id: row.id, request: JSON.parse(row.request), domain: row.domain };
});
}
async claimPostprocessJob(claimedBy: string): Promise<PostprocessJob | undefined> {
const result = await sql<{
id: number;
detail_id: number;
detail_image_path: string;
}>`
UPDATE postprocess_job
SET
claimed_by = ${claimedBy},
claimed_at = datetime('now')
WHERE id = (
SELECT id FROM postprocess_job
WHERE claimed_at IS NULL
ORDER BY created_at ASC
LIMIT 1
)
RETURNING
id,
detail_id,
(SELECT path FROM file WHERE file.id = (
SELECT image_file_id FROM detail WHERE detail.id = detail_id
)) as detail_image_path
`.execute(this.db);
const row = result.rows[0];
if (!row) return undefined;
return { id: row.id, detailId: row.detail_id, detailImagePath: row.detail_image_path };
}
async savePostprocessedFiles(
detailId: number,
thumbnailPath: string,
tilePaths: string[],
): Promise<void> {
await this.db.transaction().execute(async (trx) => {
const ensureFile = (filePath: string) =>
trx
.insertInto("file")
.values({ path: filePath })
.onConflict((oc) =>
oc.column("path").doUpdateSet({ path: (eb) => eb.ref("excluded.path") }),
)
.returning("id")
.executeTakeFirstOrThrow();
const thumbFile = await ensureFile(thumbnailPath);
await trx
.insertInto("thumbnail")
.values({ detail_id: detailId, file_id: thumbFile.id })
.onConflict((oc) =>
oc.column("detail_id").doUpdateSet({ file_id: (eb) => eb.ref("excluded.file_id") }),
)
.execute();
for (let i = 0; i < tilePaths.length; i++) {
const tileFile = await ensureFile(tilePaths[i]!);
await trx
.insertInto("tile")
.values({ detail_id: detailId, tile_index: i, file_id: tileFile.id })
.onConflict((oc) =>
oc
.columns(["detail_id", "tile_index"])
.doUpdateSet({ file_id: (eb) => eb.ref("excluded.file_id") }),
)
.execute();
}
});
}
async deletePostprocessJob(postprocessJobId: number): Promise<void> {
await this.db.deleteFrom("postprocess_job").where("id", "=", postprocessJobId).execute();
}
async saveDetail(
source: Source,
url: string | null,
title: string | null,
detailImagePath: string,
metadata: unknown,
): Promise<{ id: number; cardId: number }> {
const { uid, serialized } = await deriveUid(source);
return this.db.transaction().execute(async (trx) => {
const card = await trx
.insertInto("card")
.values({
uid: Buffer.from(uid),
source: sql<Buffer>`jsonb(${serialized})`,
url,
})
.onConflict((oc) =>
oc.column("uid").doUpdateSet({ url: (eb) => eb.ref("excluded.url") }),
)
.returning("id")
.executeTakeFirstOrThrow();
const detailImageFile = await trx
.insertInto("file")
.values({ path: detailImagePath })
.onConflict((oc) =>
oc.column("path").doUpdateSet({ path: (eb) => eb.ref("excluded.path") }),
)
.returning("id")
.executeTakeFirstOrThrow();
const detail = await trx
.insertInto("detail")
.values({
card_id: card.id,
image_file_id: detailImageFile.id,
title,
metadata: sql<Buffer>`jsonb(${JSON.stringify(metadata)})`,
})
.returning(["id", "card_id"])
.executeTakeFirstOrThrow();
await trx.insertInto("postprocess_job").values({ detail_id: detail.id }).execute();
return { id: detail.id, cardId: detail.card_id };
});
}
async saveCrawl(url: string, harPath: string, metadata: {}): Promise<number> {
const harFile = await this.db
.insertInto("file")
.values({ path: harPath })
.onConflict((oc) =>
oc.column("path").doUpdateSet({ path: (eb) => eb.ref("excluded.path") }),
)
.returning("id")
.executeTakeFirstOrThrow();
const crawl = await this.db
.insertInto("crawl")
.values({
url,
har_file_id: harFile.id,
metadata: sql<Buffer>`jsonb(${JSON.stringify(metadata)})`,
})
.returning("id")
.executeTakeFirstOrThrow();
return crawl.id;
}
async deleteCollectJob(collectJobId: number): Promise<void> {
await this.db.deleteFrom("collect_job").where("id", "=", collectJobId).execute();
}
async listDecks(): Promise<{ id: number }[]> {
return this.db.selectFrom("deck").select("id").execute();
}
async createDeck(): Promise<{ id: number }> {
return this.db
.insertInto("deck")
.defaultValues()
.returning("id")
.executeTakeFirstOrThrow();
}
async getCards(cardUids: string[]): Promise<{ cards: Card[] }> {
const hexUids = cardUids.map((uid) => encodeUid(uid).toString("hex"));
const hexUidsJson = JSON.stringify(hexUids);
const result = await sql<{
hexUid: string;
url: string | null;
title: string | null;
numTiles: number;
}>`
SELECT
uid.value AS hexUid,
card.url,
detail.title,
COUNT(tile.id) AS numTiles
FROM json_each(${hexUidsJson}) AS uid
LEFT JOIN card ON card.uid = unhex(uid.value)
LEFT JOIN detail ON detail.id = (
SELECT id FROM detail
WHERE card_id = card.id
ORDER BY id DESC
LIMIT 1
)
LEFT JOIN tile ON tile.detail_id = detail.id
GROUP BY uid.key
ORDER BY uid.key
`.execute(this.db);
return {
cards: result.rows.map((row) => ({
...row,
uid: decodeUid(Buffer.from(row.hexUid, "hex")),
})),
};
}
async getRevision(
deckId: number,
revisionId?: number,
): Promise<{ id: number; cards: Card[] } | undefined> {
const revision = revisionId
? await this.db
.selectFrom("deck")
.innerJoin("revision", "revision.deck_id", "deck.id")
.innerJoin("card_set", "card_set.id", "revision.card_set_id")
.select(["revision.id", sql<string>`json(card_set.card_ids)`.as("card_ids")])
.where("deck.id", "=", deckId)
.where("revision.id", "=", revisionId)
.executeTakeFirst()
: await this.db
.selectFrom("deck")
.leftJoin("revision", "revision.deck_id", "deck.id")
.leftJoin("card_set", "card_set.id", "revision.card_set_id")
.select(["revision.id", sql<string>`json(card_set.card_ids)`.as("card_ids")])
.where("deck.id", "=", deckId)
.orderBy("revision.id", "desc")
.limit(1)
.executeTakeFirst();
if (!revision || revision.id == null) return undefined;
const cardIds = JSON.parse(revision.card_ids) as number[];
const cards = await Promise.all(
cardIds.map(async (id) => {
const row = await sql<{
uid: Buffer;
url: string | null;
title: string | null;
numTiles: number;
}>`
SELECT
card.uid,
card.url,
detail.title,
COUNT(tile.id) AS numTiles
FROM card
LEFT JOIN detail ON detail.id = (
SELECT id FROM detail
WHERE card_id = card.id
ORDER BY id DESC
LIMIT 1
)
LEFT JOIN tile ON tile.detail_id = detail.id
WHERE card.id = ${id}
GROUP BY card.id
`
.execute(this.db)
.then((r) => r.rows[0]!);
return { ...row, uid: decodeUid(row.uid) };
}),
);
return { id: revision.id, cards };
}
async createRevision(deckId: number, cardIds: number[]): Promise<{ id: number }> {
const cardIdsJson = JSON.stringify(cardIds);
return this.db.transaction().execute(async (trx) => {
const cardSet = await trx
.insertInto("card_set")
.values({ card_ids: sql<Buffer>`jsonb(${cardIdsJson})` })
.onConflict((oc) =>
oc
.column("card_ids")
.doUpdateSet({ card_ids: (eb) => eb.ref("excluded.card_ids") }),
)
.returning("id")
.executeTakeFirstOrThrow();
return trx
.insertInto("revision")
.values({ deck_id: deckId, card_set_id: cardSet.id })
.returning("id")
.executeTakeFirstOrThrow();
});
}
async getThumbnailPath(cardUid: string): Promise<string | undefined> {
const row = await this.db
.selectFrom("card")
.innerJoin("detail", "detail.card_id", "card.id")
.innerJoin("thumbnail", "thumbnail.detail_id", "detail.id")
.innerJoin("file", "file.id", "thumbnail.file_id")
.select("file.path")
.where("card.uid", "=", encodeUid(cardUid))
.orderBy("detail.id", "desc")
.limit(1)
.executeTakeFirst();
return row?.path;
}
async getTilePath(cardUid: string, tileIndex: number): Promise<string | undefined> {
const row = await this.db
.selectFrom("card")
.innerJoin("detail", "detail.card_id", "card.id")
.innerJoin("tile", "tile.detail_id", "detail.id")
.innerJoin("file", "file.id", "tile.file_id")
.select("file.path")
.where("card.uid", "=", encodeUid(cardUid))
.where("tile.tile_index", "=", tileIndex)
.orderBy("detail.id", "desc")
.limit(1)
.executeTakeFirst();
return row?.path;
}
async createWorkspace(): Promise<{ id: number }> {
return this.db.transaction().execute(async (trx) => {
const workspace = await trx
.insertInto("workspace")
.defaultValues()
.returning("id")
.executeTakeFirstOrThrow();
const command: Command = ["spliceCards", 0, 0, 0, []];
await trx
.insertInto("operation")
.values({
workspace_id: workspace.id,
command: sql<Buffer>`jsonb(${JSON.stringify(command)})`,
})
.execute();
return workspace;
});
}
async getWorkspaceOperations(
workspaceId: number,
fromSeq: number,
excludeClient: string,
): Promise<{ seq: number; commands: Command[]; excludedCount: number } | null> {
const rows = await this.db
.selectFrom("operation")
.select([
"id",
sql<string>`json(command)`.as("command"),
sql<number>`client = ${excludeClient}`.as("is_client"),
])
.where("workspace_id", "=", workspaceId)
.where("id", ">=", fromSeq)
.orderBy("id")
.execute();
const commands: Command[] = [];
let seq: number | undefined;
let excludedCount = 0;
for (const row of rows) {
seq = row.id;
if (row.is_client) {
excludedCount++;
} else {
commands.push(JSON.parse(row.command));
}
}
return seq == null ? null : { seq, commands, excludedCount };
}
async appendWorkspaceOperations(
workspaceId: number,
client: string,
expectedSeq: number,
commands: Command[],
): Promise<number | null> {
return this.db.transaction().execute(async (trx) => {
const result = await trx
.selectFrom("operation")
.select((eb) => eb.fn.max("id").as("seq"))
.where("workspace_id", "=", workspaceId)
.executeTakeFirst();
const currentSeq = result?.seq ?? null;
if (currentSeq !== expectedSeq) return null;
let lastSeq: number = currentSeq as number;
for (const command of commands) {
const row = await trx
.insertInto("operation")
.values({
workspace_id: workspaceId,
command: sql<Buffer>`jsonb(${JSON.stringify(command)})`,
client,
})
.returning("id")
.executeTakeFirstOrThrow();
lastSeq = row.id;
}
return lastSeq;
});
}
}