From 6594299f8258fb5293ee1f6b8935ff35c2326fa9 Mon Sep 17 00:00:00 2001 From: pinks Date: Mon, 20 Nov 2023 01:33:38 +0100 Subject: [PATCH] rewrite to elysia - part 2 --- api/adminsRoute.ts | 115 +++++++++++++++++------------- api/getUser.ts | 45 ++++++++++++ api/paramsRoute.ts | 39 +++++------ api/sessionsRoute.ts | 4 +- api/statsRoute.ts | 101 +++++++++++--------------- api/usersRoute.ts | 14 ++-- api/withUser.ts | 39 ----------- api/workersRoute.ts | 136 ++++++++++++++++-------------------- app/adminStore.ts | 29 +++----- app/config.ts | 80 ++++++++++----------- app/dailyStatsStore.ts | 26 +++---- app/globalStats.ts | 23 +++--- app/userDailyStatsStore.ts | 29 ++++---- app/userStatsStore.ts | 39 ++++------- app/workerInstanceStore.ts | 46 +++++------- bot/broadcastCommand.ts | 4 +- bot/txt2imgCommand.ts | 4 +- deno.json | 6 +- deno.lock | 8 ++- ui/AdminsPage.tsx | 60 ++++++---------- ui/App.tsx | 10 +-- ui/AppHeader.tsx | 9 ++- ui/WorkersPage.tsx | 67 +++++++----------- ui/apiClient.ts | 20 ++++-- utils/Tkv.ts | 76 ++++++++++++++++++++ {app => utils}/kvMemoize.ts | 0 26 files changed, 506 insertions(+), 523 deletions(-) create mode 100644 api/getUser.ts delete mode 100644 api/withUser.ts create mode 100644 utils/Tkv.ts rename {app => utils}/kvMemoize.ts (100%) diff --git a/api/adminsRoute.ts b/api/adminsRoute.ts index 7c50005..72b04f2 100644 --- a/api/adminsRoute.ts +++ b/api/adminsRoute.ts @@ -1,26 +1,22 @@ -import { Model } from "indexed_kv"; import { info } from "std/log/mod.ts"; -import { Elysia, t } from "elysia"; +import { Elysia, Static, t } from "elysia"; import { Admin, adminSchema, adminStore } from "../app/adminStore.ts"; -import { getUser, withAdmin, withUser } from "./withUser.ts"; +import { getUser, withSessionAdmin, withSessionUser } from "./getUser.ts"; +import { TkvEntry } from "../utils/Tkv.ts"; -export type AdminData = Admin & { id: string }; +const adminDataSchema = t.Intersect([adminSchema, t.Object({ tgUserId: t.Number() })]); -const adminDataSchema = t.Object({ - id: t.String(), - tgUserId: t.Number(), - promotedBy: t.Nullable(t.String()), -}); +export type AdminData = Static; -function getAdminData(adminEntry: Model): AdminData { - return { id: adminEntry.id, ...adminEntry.value }; +function getAdminData(adminEntry: TkvEntry<["admins", number], Admin>): AdminData { + return { tgUserId: adminEntry.key[1], ...adminEntry.value }; } export const adminsRoute = new Elysia() .get( "", async () => { - const adminEntries = await adminStore.getAll(); + const adminEntries = await Array.fromAsync(adminStore.list({ prefix: ["admins"] })); const admins = adminEntries.map(getAdminData); return admins; }, @@ -30,15 +26,18 @@ export const adminsRoute = new Elysia() ) .post( "", - async ({ query, body }) => { - return withAdmin(query, async (user, adminEntry) => { + async ({ query, body, set }) => { + return withSessionAdmin({ query, set }, async (sessionUser, sessionAdminEntry) => { const newAdminUser = await getUser(body.tgUserId); - const newAdminEntry = await adminStore.create({ - tgUserId: body.tgUserId, - promotedBy: adminEntry.id, - }); - info(`User ${user.first_name} promoted user ${newAdminUser.first_name} to admin`); - return getAdminData(newAdminEntry); + const newAdminKey = ["admins", body.tgUserId] as const; + const newAdminValue = { promotedBy: sessionAdminEntry.key[1] }; + const newAdminResult = await adminStore.atomicSet(newAdminKey, null, newAdminValue); + if (!newAdminResult.ok) { + set.status = 409; + return "User is already an admin"; + } + info(`User ${sessionUser.first_name} promoted user ${newAdminUser.first_name} to admin`); + return getAdminData({ ...newAdminResult, key: newAdminKey, value: newAdminValue }); }); }, { @@ -46,62 +45,84 @@ export const adminsRoute = new Elysia() body: t.Object({ tgUserId: t.Number(), }), - response: adminDataSchema, + response: { + 200: adminDataSchema, + 401: t.Literal("Must be logged in"), + 403: t.Literal("Must be an admin"), + 409: t.Literal("User is already an admin"), + }, }, ) .post( "/promote_self", // if there are no admins, allow any user to promote themselves - async ({ query }) => { - return withUser(query, async (user) => { - const adminEntries = await adminStore.getAll(); - if (adminEntries.length === 0) { - const newAdminEntry = await adminStore.create({ - tgUserId: user.id, - promotedBy: null, - }); - info(`User ${user.first_name} promoted themselves to admin`); - return getAdminData(newAdminEntry); + async ({ query, set }) => { + return withSessionUser({ query, set }, async (sessionUser) => { + const adminEntries = await Array.fromAsync(adminStore.list({ prefix: ["admins"] })); + if (adminEntries.length !== 0) { + set.status = 409; + return "You are not allowed to promote yourself"; } - throw new Error("You are not allowed to promote yourself"); + const newAdminKey = ["admins", sessionUser.id] as const; + const newAdminValue = { promotedBy: null }; + const newAdminResult = await adminStore.set(newAdminKey, newAdminValue); + info(`User ${sessionUser.first_name} promoted themselves to admin`); + return getAdminData({ ...newAdminResult, key: newAdminKey, value: newAdminValue }); }); }, { query: t.Object({ sessionId: t.String() }), - response: adminDataSchema, + response: { + 200: adminDataSchema, + 401: t.Literal("Must be logged in"), + 403: t.Literal("Must be an admin"), + 409: t.Literal("You are not allowed to promote yourself"), + }, }, ) .get( "/:adminId", - async ({ params }) => { - const adminEntry = await adminStore.getById(params.adminId!); - if (!adminEntry) { - throw new Error("Admin not found"); + async ({ params, set }) => { + const adminEntry = await adminStore.get(["admins", Number(params.adminId)]); + if (!adminEntry.versionstamp) { + set.status = 404; + return "Admin not found"; } return getAdminData(adminEntry); }, { params: t.Object({ adminId: t.String() }), - response: adminDataSchema, + response: { + 200: adminDataSchema, + 404: t.Literal("Admin not found"), + }, }, ) .delete( "/:adminId", - async ({ params, query }) => { - return withAdmin(query, async (chat) => { - const deletedAdminEntry = await adminStore.getById(params.adminId!); - if (!deletedAdminEntry) { - throw new Error("Admin not found"); + async ({ params, query, set }) => { + return withSessionAdmin({ query, set }, async (sessionUser) => { + const deletedAdminEntry = await adminStore.get(["admins", Number(params.adminId)]); + if (!deletedAdminEntry.versionstamp) { + set.status = 404; + return "Admin not found"; } - const deletedAdminUser = await getUser(deletedAdminEntry.value.tgUserId); - await deletedAdminEntry.delete(); - info(`User ${chat.first_name} demoted user ${deletedAdminUser.first_name} from admin`); + const deletedAdminUser = await getUser(deletedAdminEntry.key[1]); + await adminStore.delete(["admins", Number(params.adminId)]); + info( + `User ${sessionUser.first_name} demoted user ${deletedAdminUser.first_name} from admin`, + ); return null; }); }, { params: t.Object({ adminId: t.String() }), query: t.Object({ sessionId: t.String() }), - response: t.Null(), + response: { + 200: t.Null(), + 401: t.Literal("Must be logged in"), + 403: t.Literal("Must be an admin"), + 404: t.Literal("Admin not found"), + }, }, ); diff --git a/api/getUser.ts b/api/getUser.ts new file mode 100644 index 0000000..0f84101 --- /dev/null +++ b/api/getUser.ts @@ -0,0 +1,45 @@ +import { Chat } from "grammy_types"; +import { Admin, adminStore } from "../app/adminStore.ts"; +import { bot } from "../bot/mod.ts"; +import { sessions } from "./sessionsRoute.ts"; +import { TkvEntry } from "../utils/Tkv.ts"; + +export async function withSessionUser( + { query, set }: { query: { sessionId: string }; set: { status?: string | number } }, + cb: (sessionUser: Chat.PrivateGetChat) => Promise, +) { + const session = sessions.get(query.sessionId); + if (!session?.userId) { + set.status = 401; + return "Must be logged in"; + } + const user = await getUser(session.userId); + return cb(user); +} + +export async function withSessionAdmin( + { query, set }: { query: { sessionId: string }; set: { status?: string | number } }, + cb: ( + sessionUser: Chat.PrivateGetChat, + sessionAdminEntry: TkvEntry<["admins", number], Admin>, + ) => Promise, +) { + const session = sessions.get(query.sessionId); + if (!session?.userId) { + set.status = 401; + return "Must be logged in"; + } + const sessionUser = await getUser(session.userId); + const sessionAdminEntry = await adminStore.get(["admins", sessionUser.id]); + if (!sessionAdminEntry.versionstamp) { + set.status = 403; + return "Must be an admin"; + } + return cb(sessionUser, sessionAdminEntry); +} + +export async function getUser(userId: number): Promise { + const chat = await bot.api.getChat(userId); + if (chat.type !== "private") throw new Error("Chat is not private"); + return chat; +} diff --git a/api/paramsRoute.ts b/api/paramsRoute.ts index 4721464..fe82bfd 100644 --- a/api/paramsRoute.ts +++ b/api/paramsRoute.ts @@ -1,46 +1,39 @@ import { Elysia, t } from "elysia"; -import { deepMerge } from "std/collections/deep_merge.ts"; import { info } from "std/log/mod.ts"; -import { getConfig, setConfig } from "../app/config.ts"; -import { omitUndef } from "../utils/omitUndef.ts"; -import { withAdmin } from "./withUser.ts"; - -const paramsSchema = t.Partial(t.Object({ - batch_size: t.Number(), - n_iter: t.Number(), - width: t.Number(), - height: t.Number(), - steps: t.Number(), - cfg_scale: t.Number(), - sampler_name: t.String(), - negative_prompt: t.String(), -})); +import { defaultParamsSchema, getConfig, setConfig } from "../app/config.ts"; +import { withSessionAdmin } from "./getUser.ts"; export const paramsRoute = new Elysia() .get( "", async () => { const config = await getConfig(); - return omitUndef(config.defaultParams); + return config.defaultParams; }, { - response: paramsSchema, + response: { + 200: defaultParamsSchema, + }, }, ) .patch( "", - async ({ query, body }) => { - return withAdmin(query, async (user) => { + async ({ query, body, set }) => { + return withSessionAdmin({ query, set }, async (user) => { const config = await getConfig(); info(`User ${user.first_name} updated default params: ${JSON.stringify(body)}`); - const defaultParams = deepMerge(config.defaultParams ?? {}, body); + const defaultParams = { ...config.defaultParams, ...body }; await setConfig({ defaultParams }); - return omitUndef(config.defaultParams); + return config.defaultParams; }); }, { query: t.Object({ sessionId: t.String() }), - body: paramsSchema, - response: paramsSchema, + body: defaultParamsSchema, + response: { + 200: defaultParamsSchema, + 401: t.Literal("Must be logged in"), + 403: t.Literal("Must be an admin"), + }, }, ); diff --git a/api/sessionsRoute.ts b/api/sessionsRoute.ts index 40ee53e..6f551fe 100644 --- a/api/sessionsRoute.ts +++ b/api/sessionsRoute.ts @@ -1,4 +1,4 @@ -import { Elysia, t } from "elysia"; +import { Elysia, NotFoundError, t } from "elysia"; import { ulid } from "ulid"; export const sessions = new Map(); @@ -29,7 +29,7 @@ export const sessionsRoute = new Elysia() const id = params.sessionId!; const session = sessions.get(id); if (!session) { - throw new Error("Session not found"); + throw new NotFoundError("Session not found"); } return { id, userId: session.userId ?? null }; }, diff --git a/api/statsRoute.ts b/api/statsRoute.ts index 2aaab7b..729fa32 100644 --- a/api/statsRoute.ts +++ b/api/statsRoute.ts @@ -1,11 +1,11 @@ import { Elysia, t } from "elysia"; import { subMinutes } from "date-fns"; -import { getDailyStats } from "../app/dailyStatsStore.ts"; +import { dailyStatsSchema, getDailyStats } from "../app/dailyStatsStore.ts"; import { generationStore } from "../app/generationStore.ts"; import { globalStats } from "../app/globalStats.ts"; -import { getUserDailyStats } from "../app/userDailyStatsStore.ts"; -import { getUserStats } from "../app/userStatsStore.ts"; -import { withAdmin } from "./withUser.ts"; +import { getUserDailyStats, userDailyStatsSchema } from "../app/userDailyStatsStore.ts"; +import { getUserStats, userStatsSchema } from "../app/userStatsStore.ts"; +import { withSessionAdmin } from "./getUser.ts"; const STATS_INTERVAL_MIN = 3; @@ -48,32 +48,25 @@ export const statsRoute = new Elysia() }; }, { - response: t.Object({ - imageCount: t.Number(), - stepCount: t.Number(), - pixelCount: t.Number(), - pixelStepCount: t.Number(), - userCount: t.Number(), - imagesPerMinute: t.Number(), - stepsPerMinute: t.Number(), - pixelsPerMinute: t.Number(), - pixelStepsPerMinute: t.Number(), - }), + response: { + 200: t.Object({ + imageCount: t.Number(), + stepCount: t.Number(), + pixelCount: t.Number(), + pixelStepCount: t.Number(), + userCount: t.Number(), + imagesPerMinute: t.Number(), + stepsPerMinute: t.Number(), + pixelsPerMinute: t.Number(), + pixelStepsPerMinute: t.Number(), + }), + }, }, ) .get( "/daily/:year/:month/:day", async ({ params }) => { - const year = Number(params.year); - const month = Number(params.month); - const day = Number(params.day); - const stats = await getDailyStats(year, month, day); - return { - imageCount: stats.imageCount, - pixelCount: stats.pixelCount, - userCount: stats.userIds.length, - timestamp: stats.timestamp, - }; + return getDailyStats(params.year, params.month, params.day); }, { params: t.Object({ @@ -81,40 +74,31 @@ export const statsRoute = new Elysia() month: t.Number(), day: t.Number(), }), - response: t.Object({ - imageCount: t.Number(), - pixelCount: t.Number(), - userCount: t.Number(), - timestamp: t.Number(), - }), + response: { + 200: dailyStatsSchema, + }, }, ) .get( "/users/:userId", async ({ params }) => { const userId = params.userId; - const stats = await getUserStats(userId); - return { - imageCount: stats.imageCount, - pixelCount: stats.pixelCount, - timestamp: stats.timestamp, - }; + // deno-lint-ignore no-unused-vars + const { tagCountMap, ...stats } = await getUserStats(userId); + return stats; }, { params: t.Object({ userId: t.Number() }), - response: t.Object({ - imageCount: t.Number(), - pixelCount: t.Number(), - timestamp: t.Number(), - }), + response: { + 200: t.Omit(userStatsSchema, ["tagCountMap"]), + }, }, ) .get( "/users/:userId/tagcount", - async ({ params, query }) => { - return withAdmin(query, async () => { - const userId = Number(params.userId); - const stats = await getUserStats(userId); + async ({ params, query, set }) => { + return withSessionAdmin({ query, set }, async () => { + const stats = await getUserStats(params.userId); return { tagCountMap: stats.tagCountMap, timestamp: stats.timestamp, @@ -124,22 +108,17 @@ export const statsRoute = new Elysia() { params: t.Object({ userId: t.Number() }), query: t.Object({ sessionId: t.String() }), - response: t.Object({ - tagCountMap: t.Record(t.String(), t.Number()), - timestamp: t.Number(), - }), + response: { + 200: t.Pick(userStatsSchema, ["tagCountMap", "timestamp"]), + 401: t.Literal("Must be logged in"), + 403: t.Literal("Must be an admin"), + }, }, ) .get( "/users/:userId/daily/:year/:month/:day", async ({ params }) => { - const { userId, year, month, day } = params; - const stats = await getUserDailyStats(userId, year, month, day); - return { - imageCount: stats.imageCount, - pixelCount: stats.pixelCount, - timestamp: stats.timestamp, - }; + return getUserDailyStats(params.userId, params.year, params.month, params.day); }, { params: t.Object({ @@ -148,10 +127,8 @@ export const statsRoute = new Elysia() month: t.Number(), day: t.Number(), }), - response: t.Object({ - imageCount: t.Number(), - pixelCount: t.Number(), - timestamp: t.Number(), - }), + response: { + 200: userDailyStatsSchema, + }, }, ); diff --git a/api/usersRoute.ts b/api/usersRoute.ts index b50c083..3d79b3e 100644 --- a/api/usersRoute.ts +++ b/api/usersRoute.ts @@ -1,7 +1,7 @@ import { Elysia, t } from "elysia"; -import { adminStore } from "../app/adminStore.ts"; +import { adminSchema, adminStore } from "../app/adminStore.ts"; import { bot } from "../bot/mod.ts"; -import { getUser } from "./withUser.ts"; +import { getUser } from "./getUser.ts"; export const usersRoute = new Elysia() .get( @@ -31,15 +31,14 @@ export const usersRoute = new Elysia() "/:userId", async ({ params }) => { const user = await getUser(Number(params.userId)); - const [adminEntry] = await adminStore.getBy("tgUserId", { value: user.id }); - const admin = adminEntry?.value; + const adminEntry = await adminStore.get(["admins", user.id]); return { id: user.id, first_name: user.first_name, last_name: user.last_name ?? null, username: user.username ?? null, bio: user.bio ?? null, - admin: admin ?? null, + admin: adminEntry.value ?? null, }; }, { @@ -50,10 +49,7 @@ export const usersRoute = new Elysia() last_name: t.Nullable(t.String()), username: t.Nullable(t.String()), bio: t.Nullable(t.String()), - admin: t.Nullable(t.Object({ - tgUserId: t.Number(), - promotedBy: t.Nullable(t.String()), - })), + admin: t.Nullable(adminSchema), }), }, ); diff --git a/api/withUser.ts b/api/withUser.ts deleted file mode 100644 index 7d3920e..0000000 --- a/api/withUser.ts +++ /dev/null @@ -1,39 +0,0 @@ -import { Chat } from "grammy_types"; -import { Model } from "indexed_kv"; -import { Admin, adminStore } from "../app/adminStore.ts"; -import { bot } from "../bot/mod.ts"; -import { sessions } from "./sessionsRoute.ts"; - -export async function withUser( - query: { sessionId: string }, - cb: (user: Chat.PrivateGetChat) => Promise, -) { - const session = sessions.get(query.sessionId); - if (!session?.userId) { - throw new Error("Must be logged in"); - } - const user = await getUser(session.userId); - return cb(user); -} - -export async function withAdmin( - query: { sessionId: string }, - cb: (user: Chat.PrivateGetChat, admin: Model) => Promise, -) { - const session = sessions.get(query.sessionId); - if (!session?.userId) { - throw new Error("Must be logged in"); - } - const user = await getUser(session.userId); - const [admin] = await adminStore.getBy("tgUserId", { value: session.userId }); - if (!admin) { - throw new Error("Must be admin"); - } - return cb(user, admin); -} - -export async function getUser(userId: number): Promise { - const chat = await bot.api.getChat(userId); - if (chat.type !== "private") throw new Error("Chat is not private"); - return chat; -} diff --git a/api/workersRoute.ts b/api/workersRoute.ts index b90b85a..b626a31 100644 --- a/api/workersRoute.ts +++ b/api/workersRoute.ts @@ -2,66 +2,38 @@ import { subMinutes } from "date-fns"; import { Model } from "indexed_kv"; import createOpenApiFetch from "openapi_fetch"; import { info } from "std/log/mod.ts"; -import { Elysia, t } from "elysia"; +import { Elysia, NotFoundError, Static, t } from "elysia"; import { activeGenerationWorkers } from "../app/generationQueue.ts"; import { generationStore } from "../app/generationStore.ts"; import * as SdApi from "../app/sdApi.ts"; import { WorkerInstance, - workerInstanceSchema as _, + workerInstanceSchema, workerInstanceStore, } from "../app/workerInstanceStore.ts"; import { getAuthHeader } from "../utils/getAuthHeader.ts"; import { omitUndef } from "../utils/omitUndef.ts"; -import { withAdmin } from "./withUser.ts"; +import { withSessionAdmin } from "./getUser.ts"; -const workerInstanceSchema = t.Object({ - name: t.Nullable(t.String()), - key: t.String(), - sdUrl: t.String(), - sdAuth: t.Nullable(t.Object({ - user: t.String(), - password: t.String(), - })), - lastOnlineTime: t.Number(), - lastError: t.Optional(t.Object({ - message: t.String(), - time: t.Number(), - })), -}); +const workerResponseSchema = t.Intersect([ + t.Object({ id: t.String() }), + t.Omit(workerInstanceSchema, ["sdUrl", "sdAuth"]), + t.Object({ + isActive: t.Boolean(), + imagesPerMinute: t.Number(), + stepsPerMinute: t.Number(), + pixelsPerMinute: t.Number(), + pixelStepsPerMinute: t.Number(), + }), +]); -export type WorkerData = { - id: string; - name: string | null; - key: string; - lastError?: { message: string; time: number }; - lastOnlineTime: number | null; - isActive: boolean; - imagesPerMinute: number; - stepsPerMinute: number; - pixelsPerMinute: number; - pixelStepsPerMinute: number; -}; +export type WorkerResponse = Static; -const workerDataSchema = t.Object({ - id: t.String(), - key: t.String(), - name: t.Nullable(t.String()), - lastError: t.Optional(t.Object({ - message: t.String(), - time: t.Number(), - })), - lastOnlineTime: t.Nullable(t.Number()), - isActive: t.Boolean(), - imagesPerMinute: t.Number(), - stepsPerMinute: t.Number(), - pixelsPerMinute: t.Number(), - pixelStepsPerMinute: t.Number(), -}); +const workerRequestSchema = t.Omit(workerInstanceSchema, ["lastOnlineTime", "lastError"]); const STATS_INTERVAL_MIN = 10; -async function getWorkerData(workerInstance: Model): Promise { +async function getWorkerResponse(workerInstance: Model): Promise { const after = subMinutes(new Date(), STATS_INTERVAL_MIN); const generations = await generationStore.getBy("workerInstanceKey", { @@ -91,7 +63,7 @@ async function getWorkerData(workerInstance: Model): Promise { const workerInstances = await workerInstanceStore.getAll(); - const workers = await Promise.all(workerInstances.map(getWorkerData)); + const workers = await Promise.all(workerInstances.map(getWorkerResponse)); return workers; }, { - response: t.Array(workerDataSchema), + response: t.Array(workerResponseSchema), }, ) .post( "", - async ({ query, body }) => { - return withAdmin(query, async (user) => { + async ({ query, body, set }) => { + return withSessionAdmin({ query, set }, async (sessionUser) => { const workerInstance = await workerInstanceStore.create(body); - info(`User ${user.first_name} created worker ${workerInstance.value.name}`); - return await getWorkerData(workerInstance); + info(`User ${sessionUser.first_name} created worker ${workerInstance.value.name}`); + return await getWorkerResponse(workerInstance); }); }, { query: t.Object({ sessionId: t.String() }), - body: t.Omit(workerInstanceSchema, ["lastOnlineTime", "lastError"]), + body: workerRequestSchema, + response: { + 200: workerResponseSchema, + 401: t.Literal("Must be logged in"), + 403: t.Literal("Must be an admin"), + }, }, ) .get( "/:workerId", async ({ params }) => { - const workerInstance = await workerInstanceStore.getById(params.workerId!); + const workerInstance = await workerInstanceStore.getById(params.workerId); if (!workerInstance) { - throw new Error("Worker not found"); + throw new NotFoundError("Worker not found"); } - return await getWorkerData(workerInstance); + return await getWorkerResponse(workerInstance); }, { params: t.Object({ workerId: t.String() }), - response: workerDataSchema, + response: { + 200: workerResponseSchema, + }, }, ) .patch( "/:workerId", - async ({ params, query, body }) => { - const workerInstance = await workerInstanceStore.getById(params.workerId!); + async ({ params, query, body, set }) => { + const workerInstance = await workerInstanceStore.getById(params.workerId); if (!workerInstance) { - throw new Error("Worker not found"); + throw new NotFoundError("Worker not found"); } - return withAdmin(query, async (user) => { + return withSessionAdmin({ query, set }, async (sessionUser) => { info( - `User ${user.first_name} updated worker ${workerInstance.value.name}: ${ + `User ${sessionUser.first_name} updated worker ${workerInstance.value.name}: ${ JSON.stringify(body) }`, ); - await workerInstance.update(omitUndef(body)); - return await getWorkerData(workerInstance); + await workerInstance.update(body); + return await getWorkerResponse(workerInstance); }); }, { params: t.Object({ workerId: t.String() }), query: t.Object({ sessionId: t.String() }), - body: t.Partial(workerInstanceSchema), + body: t.Partial(workerRequestSchema), + response: { + 200: workerResponseSchema, + 401: t.Literal("Must be logged in"), + 403: t.Literal("Must be an admin"), + }, }, ) .delete( "/:workerId", - async ({ params, query }) => { - const workerInstance = await workerInstanceStore.getById(params.workerId!); + async ({ params, query, set }) => { + const workerInstance = await workerInstanceStore.getById(params.workerId); if (!workerInstance) { throw new Error("Worker not found"); } - return withAdmin(query, async (user) => { - info(`User ${user.first_name} deleted worker ${workerInstance.value.name}`); + return withSessionAdmin({ query, set }, async (sessionUser) => { + info(`User ${sessionUser.first_name} deleted worker ${workerInstance.value.name}`); await workerInstance.delete(); return null; }); @@ -179,15 +163,19 @@ export const workersRoute = new Elysia() { params: t.Object({ workerId: t.String() }), query: t.Object({ sessionId: t.String() }), - response: t.Null(), + response: { + 200: t.Null(), + 401: t.Literal("Must be logged in"), + 403: t.Literal("Must be an admin"), + }, }, ) .get( "/:workerId/loras", async ({ params }) => { - const workerInstance = await workerInstanceStore.getById(params.workerId!); + const workerInstance = await workerInstanceStore.getById(params.workerId); if (!workerInstance) { - throw new Error("Worker not found"); + throw new NotFoundError("Worker not found"); } const sdClient = createOpenApiFetch({ baseUrl: workerInstance.value.sdUrl, @@ -218,9 +206,9 @@ export const workersRoute = new Elysia() .get( "/:workerId/models", async ({ params }) => { - const workerInstance = await workerInstanceStore.getById(params.workerId!); + const workerInstance = await workerInstanceStore.getById(params.workerId); if (!workerInstance) { - throw new Error("Worker not found"); + throw new NotFoundError("Worker not found"); } const sdClient = createOpenApiFetch({ baseUrl: workerInstance.value.sdUrl, diff --git a/app/adminStore.ts b/app/adminStore.ts index 4d85c70..2219dab 100644 --- a/app/adminStore.ts +++ b/app/adminStore.ts @@ -1,24 +1,11 @@ -import { Store } from "indexed_kv"; -import { JsonSchema, jsonType } from "t_rest/server"; +import { Static, t } from "elysia"; import { db } from "./db.ts"; +import { Tkv } from "../utils/Tkv.ts"; -export const adminSchema = { - type: "object", - properties: { - tgUserId: { type: "number" }, - promotedBy: { type: ["string", "null"] }, - }, - required: ["tgUserId", "promotedBy"], -} as const satisfies JsonSchema; - -export type Admin = jsonType; - -type AdminIndices = { - tgUserId: number; -}; - -export const adminStore = new Store(db, "adminUsers", { - indices: { - tgUserId: { getValue: (adminUser) => adminUser.tgUserId }, - }, +export const adminSchema = t.Object({ + promotedBy: t.Nullable(t.Number()), }); + +export type Admin = Static; + +export const adminStore = new Tkv<["admins", number], Admin>(db); diff --git a/app/config.ts b/app/config.ts index 652887a..1b4c3a3 100644 --- a/app/config.ts +++ b/app/config.ts @@ -1,51 +1,45 @@ +import { Static, t } from "elysia"; +import { Tkv } from "../utils/Tkv.ts"; import { db } from "./db.ts"; -import { JsonSchema, jsonType } from "t_rest/server"; -export const configSchema = { - type: "object", - properties: { - pausedReason: { type: ["string", "null"] }, - maxUserJobs: { type: "number" }, - maxJobs: { type: "number" }, - defaultParams: { - type: "object", - properties: { - batch_size: { type: "number" }, - n_iter: { type: "number" }, - width: { type: "number" }, - height: { type: "number" }, - steps: { type: "number" }, - cfg_scale: { type: "number" }, - sampler_name: { type: "string" }, - negative_prompt: { type: "string" }, - }, - }, - }, - required: ["adminUsernames", "maxUserJobs", "maxJobs", "defaultParams"], -} as const satisfies JsonSchema; +export const defaultParamsSchema = t.Partial(t.Object({ + batch_size: t.Number(), + n_iter: t.Number(), + width: t.Number(), + height: t.Number(), + steps: t.Number(), + cfg_scale: t.Number(), + sampler_name: t.String(), + negative_prompt: t.String(), +})); -export type Config = jsonType; +export type DefaultParams = Static; + +export const configSchema = t.Object({ + pausedReason: t.Nullable(t.String()), + maxUserJobs: t.Number(), + maxJobs: t.Number(), + defaultParams: defaultParamsSchema, +}); + +export type Config = Static; + +export const configStore = new Tkv<["config"], Config>(db); + +const defaultConfig: Config = { + pausedReason: null, + maxUserJobs: Infinity, + maxJobs: Infinity, + defaultParams: {}, +}; export async function getConfig(): Promise { - const configEntry = await db.get(["config"]); - const config = configEntry?.value; - return { - pausedReason: config?.pausedReason ?? null, - maxUserJobs: config?.maxUserJobs ?? Infinity, - maxJobs: config?.maxJobs ?? Infinity, - defaultParams: config?.defaultParams ?? {}, - }; + const configEntry = await configStore.get(["config"]); + return { ...defaultConfig, ...configEntry.value }; } -export async function setConfig(newConfig: Partial): Promise { - const oldConfig = await getConfig(); - const config: Config = { - pausedReason: newConfig.pausedReason === undefined - ? oldConfig.pausedReason - : newConfig.pausedReason, - maxUserJobs: newConfig.maxUserJobs ?? oldConfig.maxUserJobs, - maxJobs: newConfig.maxJobs ?? oldConfig.maxJobs, - defaultParams: newConfig.defaultParams ?? oldConfig.defaultParams, - }; - await db.set(["config"], config); +export async function setConfig(newConfig: Pick): Promise { + const configEntry = await configStore.get(["config"]); + const config = { ...defaultConfig, ...configEntry.value, ...newConfig }; + await configStore.atomicSet(["config"], configEntry.versionstamp, config); } diff --git a/app/dailyStatsStore.ts b/app/dailyStatsStore.ts index 98ab91d..bd49b93 100644 --- a/app/dailyStatsStore.ts +++ b/app/dailyStatsStore.ts @@ -1,25 +1,21 @@ import { hoursToMilliseconds, isSameDay, minutesToMilliseconds } from "date-fns"; import { UTCDateMini } from "date-fns/utc"; +import { Static, t } from "elysia"; import { info } from "std/log/mod.ts"; -import { JsonSchema, jsonType } from "t_rest/server"; import { db } from "./db.ts"; import { generationStore } from "./generationStore.ts"; -import { kvMemoize } from "./kvMemoize.ts"; +import { kvMemoize } from "../utils/kvMemoize.ts"; -export const dailyStatsSchema = { - type: "object", - properties: { - userIds: { type: "array", items: { type: "number" } }, - imageCount: { type: "number" }, - stepCount: { type: "number" }, - pixelCount: { type: "number" }, - pixelStepCount: { type: "number" }, - timestamp: { type: "number" }, - }, - required: ["userIds", "imageCount", "stepCount", "pixelCount", "pixelStepCount", "timestamp"], -} as const satisfies JsonSchema; +export const dailyStatsSchema = t.Object({ + userIds: t.Array(t.Number()), + imageCount: t.Number(), + stepCount: t.Number(), + pixelCount: t.Number(), + pixelStepCount: t.Number(), + timestamp: t.Number(), +}); -export type DailyStats = jsonType; +export type DailyStats = Static; export const getDailyStats = kvMemoize( db, diff --git a/app/globalStats.ts b/app/globalStats.ts index e03ef59..603e0c7 100644 --- a/app/globalStats.ts +++ b/app/globalStats.ts @@ -1,23 +1,16 @@ import { addDays } from "date-fns"; -import { JsonSchema, jsonType } from "t_rest/server"; import { decodeTime } from "ulid"; import { getDailyStats } from "./dailyStatsStore.ts"; import { generationStore } from "./generationStore.ts"; -export const globalStatsSchema = { - type: "object", - properties: { - userIds: { type: "array", items: { type: "number" } }, - imageCount: { type: "number" }, - stepCount: { type: "number" }, - pixelCount: { type: "number" }, - pixelStepCount: { type: "number" }, - timestamp: { type: "number" }, - }, - required: ["userIds", "imageCount", "stepCount", "pixelCount", "pixelStepCount", "timestamp"], -} as const satisfies JsonSchema; - -export type GlobalStats = jsonType; +export interface GlobalStats { + userIds: number[]; + imageCount: number; + stepCount: number; + pixelCount: number; + pixelStepCount: number; + timestamp: number; +} export const globalStats: GlobalStats = await getGlobalStats(); diff --git a/app/userDailyStatsStore.ts b/app/userDailyStatsStore.ts index 64d7446..2c88eae 100644 --- a/app/userDailyStatsStore.ts +++ b/app/userDailyStatsStore.ts @@ -1,21 +1,18 @@ -import { JsonSchema, jsonType } from "t_rest/server"; -import { kvMemoize } from "./kvMemoize.ts"; +import { Static, t } from "elysia"; +import { kvMemoize } from "../utils/kvMemoize.ts"; import { db } from "./db.ts"; import { generationStore } from "./generationStore.ts"; import { hoursToMilliseconds, isSameDay, minutesToMilliseconds } from "date-fns"; import { UTCDateMini } from "date-fns/utc"; -export const userDailyStatsSchema = { - type: "object", - properties: { - imageCount: { type: "number" }, - pixelCount: { type: "number" }, - timestamp: { type: "number" }, - }, - required: ["imageCount", "pixelCount", "timestamp"], -} as const satisfies JsonSchema; +export const userDailyStatsSchema = t.Object({ + imageCount: t.Number(), + pixelCount: t.Number(), + pixelStepCount: t.Number(), + timestamp: t.Number(), +}); -export type UserDailyStats = jsonType; +export type UserDailyStats = Static; export const getUserDailyStats = kvMemoize( db, @@ -23,6 +20,7 @@ export const getUserDailyStats = kvMemoize( async (userId: number, year: number, month: number, day: number): Promise => { let imageCount = 0; let pixelCount = 0; + let pixelStepCount = 0; for await ( const generation of generationStore.listBy("fromId", { @@ -33,11 +31,15 @@ export const getUserDailyStats = kvMemoize( ) { imageCount++; pixelCount += (generation.value.info?.width ?? 0) * (generation.value.info?.height ?? 0); + pixelStepCount += (generation.value.info?.width ?? 0) * + (generation.value.info?.height ?? 0) * + (generation.value.info?.steps ?? 0); } return { imageCount, pixelCount, + pixelStepCount, timestamp: Date.now(), }; }, @@ -51,6 +53,7 @@ export const getUserDailyStats = kvMemoize( : hoursToMilliseconds(24 * 7 + Math.random() * 24 * 7); }, // should cache if the stats are non-zero - shouldCache: (result) => result.imageCount > 0 || result.pixelCount > 0, + shouldCache: (result) => + result.imageCount > 0 || result.pixelCount > 0 || result.pixelStepCount > 0, }, ); diff --git a/app/userStatsStore.ts b/app/userStatsStore.ts index e3b9aa8..0d5fb64 100644 --- a/app/userStatsStore.ts +++ b/app/userStatsStore.ts @@ -1,38 +1,23 @@ import { minutesToMilliseconds } from "date-fns"; import { Store } from "indexed_kv"; import { info } from "std/log/mod.ts"; -import { JsonSchema, jsonType } from "t_rest/server"; +import { Static, t } from "elysia"; import { db } from "./db.ts"; import { generationStore } from "./generationStore.ts"; -import { kvMemoize } from "./kvMemoize.ts"; +import { kvMemoize } from "../utils/kvMemoize.ts"; import { sortBy } from "std/collections/sort_by.ts"; -export const userStatsSchema = { - type: "object", - properties: { - userId: { type: "number" }, - imageCount: { type: "number" }, - stepCount: { type: "number" }, - pixelCount: { type: "number" }, - pixelStepCount: { type: "number" }, - tagCountMap: { - type: "object", - additionalProperties: { type: "number" }, - }, - timestamp: { type: "number" }, - }, - required: [ - "userId", - "imageCount", - "stepCount", - "pixelCount", - "pixelStepCount", - "tagCountMap", - "timestamp", - ], -} as const satisfies JsonSchema; +export const userStatsSchema = t.Object({ + userId: t.Number(), + imageCount: t.Number(), + stepCount: t.Number(), + pixelCount: t.Number(), + pixelStepCount: t.Number(), + tagCountMap: t.Record(t.String(), t.Number()), + timestamp: t.Number(), +}); -export type UserStats = jsonType; +export type UserStats = Static; type UserStatsIndices = { userId: number; diff --git a/app/workerInstanceStore.ts b/app/workerInstanceStore.ts index 06325de..144eeca 100644 --- a/app/workerInstanceStore.ts +++ b/app/workerInstanceStore.ts @@ -1,37 +1,23 @@ import { Store } from "indexed_kv"; -import { JsonSchema, jsonType } from "t_rest/server"; +import { Static, t } from "elysia"; import { db } from "./db.ts"; -export const workerInstanceSchema = { - type: "object", - properties: { - // used for counting stats - key: { type: "string" }, - // used for display - name: { type: ["string", "null"] }, - sdUrl: { type: "string" }, - sdAuth: { - type: ["object", "null"], - properties: { - user: { type: "string" }, - password: { type: "string" }, - }, - required: ["user", "password"], - }, - lastOnlineTime: { type: "number" }, - lastError: { - type: "object", - properties: { - message: { type: "string" }, - time: { type: "number" }, - }, - required: ["message", "time"], - }, - }, - required: ["key", "name", "sdUrl", "sdAuth"], -} as const satisfies JsonSchema; +export const workerInstanceSchema = t.Object({ + key: t.String(), + name: t.Nullable(t.String()), + sdUrl: t.String(), + sdAuth: t.Nullable(t.Object({ + user: t.String(), + password: t.String(), + })), + lastOnlineTime: t.Optional(t.Number()), + lastError: t.Optional(t.Object({ + message: t.String(), + time: t.Number(), + })), +}); -export type WorkerInstance = jsonType; +export type WorkerInstance = Static; export const workerInstanceStore = new Store(db, "workerInstances", { indices: {}, diff --git a/bot/broadcastCommand.ts b/bot/broadcastCommand.ts index d1d23cf..5a61182 100644 --- a/bot/broadcastCommand.ts +++ b/bot/broadcastCommand.ts @@ -12,9 +12,9 @@ export async function broadcastCommand(ctx: CommandContext) { return ctx.reply("I don't know who you are."); } - const [admin] = await adminStore.getBy("tgUserId", { value: ctx.from.id }); + const adminEntry = await adminStore.get(["admins", ctx.from.id]); - if (!admin) { + if (!adminEntry.versionstamp) { return ctx.reply("Only a bot admin can use this command."); } diff --git a/bot/txt2imgCommand.ts b/bot/txt2imgCommand.ts index b09847a..0a3087a 100644 --- a/bot/txt2imgCommand.ts +++ b/bot/txt2imgCommand.ts @@ -39,9 +39,9 @@ async function txt2img(ctx: ErisContext, match: string, includeRepliedTo: boolea return; } - const [admin] = await adminStore.getBy("tgUserId", { value: ctx.from.id }); + const adminEntry = await adminStore.get(["admins", ctx.from.id]); - if (admin) { + if (adminEntry.versionstamp) { priority = 1; } else { const jobs = await generationQueue.getAllJobs(); diff --git a/deno.json b/deno.json index 06407c6..035bb4e 100644 --- a/deno.json +++ b/deno.json @@ -12,8 +12,8 @@ "date-fns": "https://cdn.skypack.dev/date-fns@2.30.0?dts", "date-fns/utc": "https://cdn.skypack.dev/@date-fns/utc@1.1.0?dts", "elysia": "https://esm.sh/elysia@0.7.21?dev", - "elysia/eden": "https://esm.sh/@elysiajs/eden@0.7.4?dev", - "elysia/swagger": "https://esm.sh/@elysiajs/swagger@0.7.4?dev", + "elysia/eden": "https://esm.sh/@elysiajs/eden@0.7.4?external=elysia&dev", + "elysia/swagger": "https://esm.sh/@elysiajs/swagger@0.7.4?external=elysia&dev", "exifreader": "https://esm.sh/exifreader@4.14.1", "file_type": "https://esm.sh/file-type@18.5.0", "grammy": "https://lib.deno.dev/x/grammy@1/mod.ts", @@ -43,8 +43,6 @@ "std/path/": "https://deno.land/std@0.204.0/path/", "swr": "https://esm.sh/swr@2.2.4?external=react&dev", "swr/mutation": "https://esm.sh/swr@2.2.4/mutation?external=react&dev", - "t_rest/client": "https://esm.sh/ty-rest@0.4.1/client", - "t_rest/server": "https://esm.sh/ty-rest@0.4.1/server", "twind/core": "https://esm.sh/@twind/core@1.1.3", "twind/preset-tailwind": "https://esm.sh/@twind/preset-tailwind@1.1.4", "ulid": "https://deno.land/x/ulid@v0.3.0/mod.ts" diff --git a/deno.lock b/deno.lock index c0fe17e..b82d793 100644 --- a/deno.lock +++ b/deno.lock @@ -299,7 +299,9 @@ "https://deno.land/x/swc@0.2.1/lib/deno_swc.generated.js": "a112eb5a4871bcbd5b660d3fd9d7afdd5cf2bb9e135eb9b04aceaa24d16481a0", "https://deno.land/x/swc@0.2.1/mod.ts": "9e03f5daf4c2a25208e34e19ce3e997d2e23a5a11ee8c8ed58023b0e3626073f", "https://deno.land/x/ulid@v0.3.0/mod.ts": "f7ff065b66abd485051fc68af23becef6ccc7e81f7774d7fcfd894a4b2da1984", + "https://esm.sh/@elysiajs/eden@0.7.4?external=elysia&dev": "28d477942e36cdeb3e57791a3acc1f52dd16378b1e639693897b21e98d7295e3", "https://esm.sh/@elysiajs/swagger@0.7.4?dev": "78520881a756f6c3a69ccc1f306bc32e258c6e0e27f971f68231763fbafd303e", + "https://esm.sh/@elysiajs/swagger@0.7.4?external=elysia&dev": "9732dabca93af3dee2ed1781edb740f6e6bdbe88167cc6a03fa42bef3aadd315", "https://esm.sh/@grammyjs/types@2.0.0": "5e5d1ae9f014cb64f0af36fb91f9d35414d670049b77367818a70cd62092b2ac", "https://esm.sh/@grammyjs/types@2.12.1": "eebe448d3bf3d4fdaeacee50e31b9e3947ce965d2591f1036e5c0273cb7aec36", "https://esm.sh/@twind/core@1.1.3": "bf3e64c13de95b061a721831bf2aed322f6e7335848b06d805168c3212686c1d", @@ -378,6 +380,10 @@ "https://esm.sh/v133/ty-rest@0.4.1/denonext/server.js": "00be1165ac96313b077629556a1587d4662f4e23bb0e815b945cc95cd3582370", "https://esm.sh/v133/use-local-storage@3.0.0/X-ZS9yZWFjdA/denonext/use-local-storage.development.mjs": "1fdc00893fe7dac56e95e2817e05d413b674f5cb5a1c6afd8994e25c9e2a56c8", "https://esm.sh/v133/use-sync-external-store@1.2.0/X-ZS9yZWFjdA/denonext/shim.development.js": "5388baf48494f5abe76f8a4a30810c48e828b52f1298826aa9a3f3378e2b533f", - "https://esm.sh/v133/util-deprecate@1.0.2/denonext/util-deprecate.mjs": "f69f67cf655c38428b0934e0f7c865c055834a87cc3866b629d6b2beb21005e9" + "https://esm.sh/v133/util-deprecate@1.0.2/denonext/util-deprecate.mjs": "f69f67cf655c38428b0934e0f7c865c055834a87cc3866b629d6b2beb21005e9", + "https://esm.sh/v134/@elysiajs/eden@0.7.4/X-ZS9lbHlzaWE/denonext/eden.development.mjs": "c79c9f105d2b2062b882240272faa81cab1cbdde1113290d2d56c15ed9479986", + "https://esm.sh/v134/@elysiajs/swagger@0.7.4/X-ZS9lbHlzaWE/denonext/swagger.development.mjs": "4aa74e7b8108e9bacd7830c374ef24ee1b4100d65eef045bf6802ceaa3635554", + "https://esm.sh/v134/@sinclair/typebox@0.31.26/X-ZS9lbHlzaWE/denonext/typebox.development.mjs": "5f24db9ca594ccb61f2ae58df80a1f7690d21e117772e7b00babdb290c0747df", + "https://esm.sh/v134/lodash.clonedeep@4.5.0/X-ZS9lbHlzaWE/denonext/lodash.clonedeep.development.mjs": "f44c863cb41bcd2714103a36d97e8eb8268f56070b5cd56dcfae26c556e6622e" } } diff --git a/ui/AdminsPage.tsx b/ui/AdminsPage.tsx index 295d2c3..abf519a 100644 --- a/ui/AdminsPage.tsx +++ b/ui/AdminsPage.tsx @@ -32,15 +32,10 @@ export function AdminsPage(props: { sessionId: string | null }) {