refactor: add generic function getRowCount

tippfehlr · Apr 13, 2024 · d31e0b3 · d31e0b3
1 parent 952424c
commit d31e0b3
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 64 deletions.
diff --git 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<number> {
+export async function getRowCount(table: TableExpression<DB, keyof DB>): Promise<number> {
   return (
     await db
-      .selectFrom('usersHashed')
-      .select(eb => eb.fn.countAll().as('count'))
-      .executeTakeFirstOrThrow()
-  ).count as number;
-}
-
-export async function getNewUserCount(): Promise<number> {
-  return (
-    await db
-      .selectFrom('users')
+      .selectFrom(table)
       .select(eb => eb.fn.countAll().as('count'))
       .executeTakeFirstOrThrow()
   ).count as number;
 }
 
 export async function getUserCount(): Promise<number> {
-  return (await getOldUserCount()) + (await getNewUserCount());
-}
-
-export async function getActivityRoleCount(): Promise<number> {
-  return (
-    await db
-      .selectFrom('activityRoles')
-      .select(eb => eb.fn.countAll().as('count'))
-      .executeTakeFirstOrThrow()
-  ).count as number;
-}
-
-export async function getStatusRoleCount(): Promise<number> {
-  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<number> {
@@ -182,27 +163,11 @@ export async function getPermRoleCount(): Promise<number> {
 }
 
 export async function getRolesCount(): Promise<number> {
-  return (await getActivityRoleCount()) + (await getStatusRoleCount());
-}
-
-export async function getOldActiveTemporaryRolesCount(): Promise<number> {
-  return (
-    await db
-      .selectFrom('activeTemporaryRolesHashed')
-      .select(eb => eb.fn.countAll().as('count'))
-      .executeTakeFirstOrThrow()
-  ).count as number;
-}
-
-export async function getNewActiveTemporaryRolesCount(): Promise<number> {
-  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<number> {
-  return (await getOldActiveTemporaryRolesCount()) + (await getNewActiveTemporaryRolesCount());
+  return (
+    (await getRowCount('activeTemporaryRolesHashed')) + (await getRowCount('activeTemporaryRoles'))
+  );
 }
diff --git 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()),
       );