diff --git a/src/modules/db.ts b/src/modules/db.ts index 8449b22..6163aef 100644 --- a/src/modules/db.ts +++ b/src/modules/db.ts @@ -2,13 +2,21 @@ import fs from 'fs'; import { createHash } from 'crypto'; import { CommandInteraction, StringSelectMenuInteraction } from 'discord.js'; import { Pool } from 'pg'; -import { FileMigrationProvider, Kysely, Migrator, PostgresDialect, Selectable } from 'kysely'; +import { + FileMigrationProvider, + Kysely, + Migrator, + PostgresDialect, + SelectQueryBuilder, + Selectable, +} from 'kysely'; import path from 'path'; import { locales, log } from './messages'; import { DB, ActivityRoles, Guilds, StatusRoles, Users } from './db.types'; import config from './config'; import { writeIntPoint } from './metrics'; +import { ExtractTableAlias, TableExpression } from 'kysely/dist/cjs/parser/table-parser'; const pool = new Pool({ connectionString: config.DATABASE_URL, max: 10 }); const dialect = new PostgresDialect({ pool }); @@ -121,44 +129,17 @@ export async function addActivity(guildID: string, activityName: string) { .execute(); } -export async function getOldUserCount(): Promise { +export async function getRowCount(table: TableExpression): Promise { return ( await db - .selectFrom('usersHashed') - .select(eb => eb.fn.countAll().as('count')) - .executeTakeFirstOrThrow() - ).count as number; -} - -export async function getNewUserCount(): Promise { - return ( - await db - .selectFrom('users') + .selectFrom(table) .select(eb => eb.fn.countAll().as('count')) .executeTakeFirstOrThrow() ).count as number; } export async function getUserCount(): Promise { - return (await getOldUserCount()) + (await getNewUserCount()); -} - -export async function getActivityRoleCount(): Promise { - return ( - await db - .selectFrom('activityRoles') - .select(eb => eb.fn.countAll().as('count')) - .executeTakeFirstOrThrow() - ).count as number; -} - -export async function getStatusRoleCount(): Promise { - return ( - await db - .selectFrom('statusRoles') - .select(eb => eb.fn.countAll().as('count')) - .executeTakeFirstOrThrow() - ).count as number; + return (await getRowCount('users')) + (await getRowCount('usersHashed')); } export async function getTempRoleCount(): Promise { @@ -182,27 +163,11 @@ export async function getPermRoleCount(): Promise { } export async function getRolesCount(): Promise { - return (await getActivityRoleCount()) + (await getStatusRoleCount()); -} - -export async function getOldActiveTemporaryRolesCount(): Promise { - return ( - await db - .selectFrom('activeTemporaryRolesHashed') - .select(eb => eb.fn.countAll().as('count')) - .executeTakeFirstOrThrow() - ).count as number; -} - -export async function getNewActiveTemporaryRolesCount(): Promise { - return ( - await db - .selectFrom('activeTemporaryRoles') - .select(eb => eb.fn.countAll().as('count')) - .executeTakeFirstOrThrow() - ).count as number; + return (await getRowCount('activityRoles')) + (await getRowCount('statusRoles')); } export async function getActiveTemporaryRolesCount(): Promise { - return (await getOldActiveTemporaryRolesCount()) + (await getNewActiveTemporaryRolesCount()); + return ( + (await getRowCount('activeTemporaryRolesHashed')) + (await getRowCount('activeTemporaryRoles')) + ); } diff --git a/src/modules/metrics.ts b/src/modules/metrics.ts index c15f284..cc663b3 100644 --- a/src/modules/metrics.ts +++ b/src/modules/metrics.ts @@ -9,21 +9,24 @@ import { stats as botStats, resetStats as resetBotStats } from './bot'; import { getUserCount, getRolesCount, + getActiveTemporaryRolesCount, getTempRoleCount, getPermRoleCount, - getStatusRoleCount, - getOldUserCount, - getNewUserCount, - getOldActiveTemporaryRolesCount, - getNewActiveTemporaryRolesCount, - getActiveTemporaryRolesCount, + getRowCount, } from './db'; -export let client: InfluxDB; +let client: InfluxDB; export let writeApi: WriteApi; export function writeIntPoint(name: string, fieldName: string, value: number) { - if (writeApi) writeApi.writePoint(new Point(name).intField(fieldName, value)); + if (writeApi) + try { + writeApi.writePoint(new Point(name).intField(fieldName, value)); + } catch (err: any) { + if (err.message !== 'writeApi: already closed!') { + log.error(err); + } + } } export async function configureInfluxDB() { @@ -43,18 +46,18 @@ export async function configureInfluxDB() { writeIntPoint('roles', 'roles_count', await getRolesCount()); writeIntPoint('roles', 'temporary_roles_count', await getTempRoleCount()); writeIntPoint('roles', 'permanent_roles_count', await getPermRoleCount()); - writeIntPoint('roles', 'status_roles_count', await getStatusRoleCount()); + writeIntPoint('roles', 'status_roles_count', await getRowCount('statusRoles')); writeApi.writePoint( new Point('users') .intField('users_cache_total', discordClient.users.cache.size) - .intField('old_users_count', await getOldUserCount()) - .intField('new_users_count', await getNewUserCount()) + .intField('old_users_count', await getRowCount('usersHashed')) + .intField('new_users_count', await getRowCount('users')) .intField('users_db_total', await getUserCount()), ); writeApi.writePoint( new Point('activeTemporaryRoles') - .intField('old', await getOldActiveTemporaryRolesCount()) - .intField('new', await getNewActiveTemporaryRolesCount()) + .intField('old', await getRowCount('activeTemporaryRolesHashed')) + .intField('new', await getRowCount('activeTemporaryRoles')) .intField('total', await getActiveTemporaryRolesCount()), );