feat(core,schemas): add daily token usage table to ease the usage calculation (#5148)

packages/core/src/event-listeners/access-token.test.ts

@@ -1,25 +0,0 @@
-import { MockQueries } from '#src/test-utils/tenant.js';
-
-import { accessTokenIssuedListener } from './access-token.js';
-
-const { jest } = import.meta;
-const insertActiveUser = jest.fn();
-const queries = new MockQueries({
-  dailyActiveUsers: { insertActiveUser },
-});
-
-describe('accessTokenIssuedListener()', () => {
-  afterEach(() => {
-    insertActiveUser.mockClear();
-  });
-
-  it('should call insertActiveUser if accountId exists', async () => {
-    await accessTokenIssuedListener({ accountId: 'accountId' }, queries);
-    expect(insertActiveUser).toHaveBeenCalled();
-  });
-
-  it('should not call insertActiveUser if no accountId', async () => {
-    await accessTokenIssuedListener({ accountId: '' }, queries);
-    expect(insertActiveUser).not.toHaveBeenCalled();
-  });
-});

packages/core/src/event-listeners/index.ts

@@ -3,27 +3,43 @@ import type Provider from 'oidc-provider';
 import type Queries from '#src/tenants/Queries.js';
 import { consoleLog } from '#src/utils/console.js';
 
-import { accessTokenIssuedListener } from './access-token.js';
 import { grantListener, grantRevocationListener } from './grant.js';
 import { interactionEndedListener, interactionStartedListener } from './interaction.js';
+import { recordActiveUsers } from './record-active-users.js';
 
 /**
  * @see {@link https://github.com/panva/node-oidc-provider/blob/v7.x/docs/README.md#im-getting-a-client-authentication-failed-error-with-no-details Getting auth error with no details?}
  * @see {@link https://github.com/panva/node-oidc-provider/blob/v7.x/docs/events.md OIDC Provider events}
  */
 export const addOidcEventListeners = (provider: Provider, queries: Queries) => {
+  const { recordTokenUsage } = queries.dailyTokenUsage;
+  const countTokenUsage = async () => recordTokenUsage(new Date());
+
   provider.addListener('grant.success', grantListener);
   provider.addListener('grant.error', grantListener);
   provider.addListener('grant.revoked', grantRevocationListener);
   provider.addListener('access_token.issued', async (token) => {
-    return accessTokenIssuedListener(token, queries);
+    return recordActiveUsers(token, queries);
   });
   provider.addListener('access_token.saved', async (token) => {
-    return accessTokenIssuedListener(token, queries);
+    return recordActiveUsers(token, queries);
   });
   provider.addListener('interaction.started', interactionStartedListener);
   provider.addListener('interaction.ended', interactionEndedListener);
   provider.addListener('server_error', (_, error) => {
     consoleLog.error('OIDC Provider server_error:', error);
   });
+
+  // Record token usage.
+  for (const event of [
+    'access_token.saved',
+    'access_token.issued',
+    'client_credentials.saved',
+    'client_credentials.issued',
+    'initial_access_token.saved',
+    'registration_access_token.saved',
+    'refresh_token.saved',
+  ]) {
+    provider.addListener(event, countTokenUsage);
+  }
 };

packages/core/src/event-listeners/record-active-users.test.ts

@@ -0,0 +1,25 @@
+import { MockQueries } from '#src/test-utils/tenant.js';
+
+import { recordActiveUsers } from './record-active-users.js';
+
+const { jest } = import.meta;
+const insertActiveUser = jest.fn();
+const queries = new MockQueries({
+  dailyActiveUsers: { insertActiveUser },
+});
+
+describe('recordActiveUsers()', () => {
+  afterEach(() => {
+    insertActiveUser.mockClear();
+  });
+
+  it('should call insertActiveUser if accountId exists, should always call recordTokenUsage', async () => {
+    await recordActiveUsers({ accountId: 'accountId' }, queries);
+    expect(insertActiveUser).toHaveBeenCalled();
+  });
+
+  it('should not call insertActiveUser if no accountId, should always call recordTokenUsage', async () => {
+    await recordActiveUsers({ accountId: '' }, queries);
+    expect(insertActiveUser).not.toHaveBeenCalled();
+  });
+});

packages/core/src/event-listeners/record-active-users.ts

@@ -1,12 +1,9 @@
 import { generateStandardId } from '@logto/shared';
 
-import { getUtcStartOfToday } from '#src/oidc/utils.js';
+import { getUtcStartOfTheDay } from '#src/oidc/utils.js';
 import type Queries from '#src/tenants/Queries.js';
 
-export const accessTokenIssuedListener = async (
-  accessToken: { accountId?: string },
-  queries: Queries
-) => {
+export const recordActiveUsers = async (accessToken: { accountId?: string }, queries: Queries) => {
   const { accountId } = accessToken;
   const { insertActiveUser } = queries.dailyActiveUsers;
 
@@ -19,6 +16,6 @@ export const accessTokenIssuedListener = async (
   await insertActiveUser({
     id: generateStandardId(),
     userId: accountId,
-    date: getUtcStartOfToday().getTime(),
+    date: getUtcStartOfTheDay(new Date()).getTime(),
   });
 };

packages/core/src/oidc/utils.ts

@@ -64,8 +64,8 @@ export const isOriginAllowed = (
   return [...corsAllowedOrigins, ...redirectUriOrigins].includes(origin);
 };
 
-export const getUtcStartOfToday = () => {
-  const now = new Date();
-
-  return new Date(Date.UTC(now.getUTCFullYear(), now.getUTCMonth(), now.getUTCDate(), 0, 0, 0, 0));
+export const getUtcStartOfTheDay = (date: Date) => {
+  return new Date(
+    Date.UTC(date.getUTCFullYear(), date.getUTCMonth(), date.getUTCDate(), 0, 0, 0, 0)
+  );
 };

packages/core/src/queries/daily-token-usage.ts

@@ -0,0 +1,46 @@
+import { DailyTokenUsage } from '@logto/schemas';
+import { convertToIdentifiers, generateStandardId } from '@logto/shared';
+import type { CommonQueryMethods } from 'slonik';
+import { sql } from 'slonik';
+
+import { getUtcStartOfTheDay } from '#src/oidc/utils.js';
+
+const { table, fields } = convertToIdentifiers(DailyTokenUsage);
+const { fields: fieldsWithPrefix } = convertToIdentifiers(DailyTokenUsage, true);
+
+export const createDailyTokenUsageQueries = (pool: CommonQueryMethods) => {
+  /**
+   * Record the token usage of the current date.
+   *
+   * @param date The current date.
+   * @returns The updated token usage of the current date.
+   */
+  /**
+   * We opted to manually write a query instead of using existing query building methods,
+   * as the latter couldn't support the complexity of this specific query logic.
+   *
+   * If we were to use the pre-built query methods, completing this operation would
+   * require two database requests:
+   * 1. to request the record
+   * 2. to update it if the record exists, or insert a new one if it doesn’t
+   *
+   * The approach we used allows us to accomplish the task within a single database query.
+   */
+  const recordTokenUsage = async (date: Date) =>
+    // Insert a new record if not exists (with usage to be 1, since this
+    // should be the first token use of the day), otherwise increment the usage by 1.
+    pool.one<DailyTokenUsage>(sql`
+      insert into ${table} (${fields.id}, ${fields.date}, ${fields.usage})
+      values (${generateStandardId()}, to_timestamp(${getUtcStartOfTheDay(
+        date
+      ).getTime()}::double precision / 1000), 1)
+      on conflict (${fields.date}, ${fields.tenantId}) do update set ${fields.usage} = ${
+        fieldsWithPrefix.usage
+      } + 1
+      returning ${sql.join(Object.values(fields), sql`, `)}
+    `);
+
+  return {
+    recordTokenUsage,
+  };
+};

packages/core/src/tenants/Queries.ts

@@ -7,6 +7,7 @@ import { createApplicationsRolesQueries } from '#src/queries/applications-roles.
 import { createConnectorQueries } from '#src/queries/connector.js';
 import { createCustomPhraseQueries } from '#src/queries/custom-phrase.js';
 import { createDailyActiveUsersQueries } from '#src/queries/daily-active-user.js';
+import { createDailyTokenUsageQueries } from '#src/queries/daily-token-usage.js';
 import { createDomainsQueries } from '#src/queries/domains.js';
 import { createHooksQueries } from '#src/queries/hooks.js';
 import { createLogQueries } from '#src/queries/log.js';
@@ -46,6 +47,7 @@ export default class Queries {
   hooks = createHooksQueries(this.pool);
   domains = createDomainsQueries(this.pool);
   dailyActiveUsers = createDailyActiveUsersQueries(this.pool);
+  dailyTokenUsage = createDailyTokenUsageQueries(this.pool);
   organizations = new OrganizationQueries(this.pool);
   ssoConnectors = new SsoConnectorQueries(this.pool);
   userSsoIdentities = new UserSsoIdentityQueries(this.pool);

packages/schemas/alterations/next-1703229996-daily-token-usage.ts

@@ -0,0 +1,62 @@
+import { type CommonQueryMethods, sql } from 'slonik';
+
+import type { AlterationScript } from '../lib/types/alteration.js';
+
+const getId = (value: string) => sql.identifier([value]);
+
+const getDatabaseName = async (pool: CommonQueryMethods) => {
+  const { currentDatabase } = await pool.one<{ currentDatabase: string }>(sql`
+    select current_database();
+  `);
+
+  return currentDatabase.replaceAll('-', '_');
+};
+
+const enableRls = async (pool: CommonQueryMethods, database: string, table: string) => {
+  const baseRoleId = sql.identifier([`logto_tenant_${database}`]);
+
+  await pool.query(sql`
+    create trigger set_tenant_id before insert on ${sql.identifier([table])}
+      for each row execute procedure set_tenant_id();
+
+    alter table ${sql.identifier([table])} enable row level security;
+
+    create policy ${sql.identifier([`${table}_tenant_id`])} on ${sql.identifier([table])}
+      as restrictive
+      using (tenant_id = (select id from tenants where db_user = current_user));
+
+    create policy ${sql.identifier([`${table}_modification`])} on ${sql.identifier([table])}
+      using (true);
+
+    grant select, insert, update, delete on ${sql.identifier([table])} to ${baseRoleId};
+  `);
+};
+
+const alteration: AlterationScript = {
+  up: async (pool) => {
+    const database = await getDatabaseName(pool);
+
+    await pool.query(sql`
+      create table daily_token_usage (
+        id varchar(21) not null,
+        tenant_id varchar(21) not null
+          references tenants (id) on update cascade on delete cascade,
+        usage bigint not null default(0),
+        date timestamptz not null,
+        primary key (id)
+      );
+
+      create unique index daily_token_usage__date
+        on daily_token_usage (tenant_id, date);
+    `);
+
+    await enableRls(pool, database, 'daily_token_usage');
+  },
+  down: async (pool) => {
+    await pool.query(sql`
+      drop table daily_token_usage;
+    `);
+  },
+};
+
+export default alteration;

packages/schemas/tables/daily_token_usage.sql

@@ -0,0 +1,11 @@
+create table daily_token_usage (
+  id varchar(21) not null,
+  tenant_id varchar(21) not null
+    references tenants (id) on update cascade on delete cascade,
+  usage bigint not null default(0),
+  date timestamptz not null,
+  primary key (id)
+);
+
+create unique index daily_token_usage__date
+  on daily_token_usage (tenant_id, date);