From d5885160cc14f696bfb438ba3c5967a4cf4c029a Mon Sep 17 00:00:00 2001 From: wangsijie Date: Thu, 20 Jul 2023 16:29:39 +0800 Subject: [PATCH] refactor(core): create quota library (#4185) --- packages/core/src/database/row-count.ts | 9 +- packages/core/src/libraries/quota.test.ts | 44 +++++++++ packages/core/src/libraries/quota.ts | 91 +++++++++++++++++ .../src/middleware/koa-quota-guard.test.ts | 97 ------------------- .../core/src/middleware/koa-quota-guard.ts | 64 +----------- packages/core/src/queries/application.ts | 14 ++- packages/core/src/routes/application.test.ts | 52 +++++----- packages/core/src/routes/application.ts | 18 +++- packages/core/src/routes/resource.test.ts | 2 + packages/core/src/routes/resource.ts | 2 + packages/core/src/tenants/Libraries.ts | 6 +- packages/core/src/tenants/Tenant.ts | 2 +- packages/core/src/test-utils/quota.ts | 9 ++ packages/core/src/test-utils/tenant.ts | 2 +- 14 files changed, 222 insertions(+), 190 deletions(-) create mode 100644 packages/core/src/libraries/quota.test.ts create mode 100644 packages/core/src/libraries/quota.ts delete mode 100644 packages/core/src/middleware/koa-quota-guard.test.ts create mode 100644 packages/core/src/test-utils/quota.ts diff --git a/packages/core/src/database/row-count.ts b/packages/core/src/database/row-count.ts index eed209802..3d3c43921 100644 --- a/packages/core/src/database/row-count.ts +++ b/packages/core/src/database/row-count.ts @@ -2,8 +2,13 @@ import type { CommonQueryMethods, IdentifierSqlToken } from 'slonik'; import { sql } from 'slonik'; export const getTotalRowCountWithPool = - (pool: CommonQueryMethods) => async (table: IdentifierSqlToken) => - pool.one<{ count: number }>(sql` + (pool: CommonQueryMethods) => async (table: IdentifierSqlToken) => { + // Postgres returns a biging for count(*), which is then converted to a string by query library. + // We need to convert it to a number. + const { count } = await pool.one<{ count: string }>(sql` select count(*) from ${table} `); + + return { count: Number(count) }; + }; diff --git a/packages/core/src/libraries/quota.test.ts b/packages/core/src/libraries/quota.test.ts new file mode 100644 index 000000000..f03522548 --- /dev/null +++ b/packages/core/src/libraries/quota.test.ts @@ -0,0 +1,44 @@ +import { createMockUtils } from '@logto/shared/esm'; + +import { mockFreePlan } from '#src/__mocks__/subscription.js'; +import { createMockCloudConnectionLibrary } from '#src/test-utils/cloud-connection.js'; + +const { jest } = import.meta; +const { mockEsmWithActual } = createMockUtils(jest); + +const { getTenantSubscriptionPlan } = await mockEsmWithActual( + '#src/utils/subscription/index.js', + () => ({ + getTenantSubscriptionPlan: jest.fn().mockResolvedValue(mockFreePlan), + }) +); + +const cloudConnection = createMockCloudConnectionLibrary(); + +const { MockQueries } = await import('#src/test-utils/tenant.js'); +const { createQuotaLibrary } = await import('./quota.js'); + +const countNonM2mApplications = jest.fn(); +const queries = new MockQueries({ + applications: { countNonM2mApplications }, +}); + +describe('guardKey()', () => { + afterEach(() => { + getTenantSubscriptionPlan.mockClear(); + }); + + const { guardKey } = createQuotaLibrary(queries, cloudConnection); + + it('should pass when limit is not exeeded', async () => { + countNonM2mApplications.mockResolvedValueOnce(0); + + await expect(guardKey('applicationsLimit')).resolves.not.toThrow(); + }); + + it('should throw when limit is exeeded', async () => { + countNonM2mApplications.mockResolvedValueOnce(mockFreePlan.quota.applicationsLimit); + + await expect(guardKey('applicationsLimit')).rejects.toThrow(); + }); +}); diff --git a/packages/core/src/libraries/quota.ts b/packages/core/src/libraries/quota.ts new file mode 100644 index 000000000..456e80f4f --- /dev/null +++ b/packages/core/src/libraries/quota.ts @@ -0,0 +1,91 @@ +import { EnvSet } from '#src/env-set/index.js'; +import RequestError from '#src/errors/RequestError/index.js'; +import type Queries from '#src/tenants/Queries.js'; +import assertThat from '#src/utils/assert-that.js'; +import { getTenantSubscriptionPlan } from '#src/utils/subscription/index.js'; +import { type FeatureQuota } from '#src/utils/subscription/types.js'; + +import { type CloudConnectionLibrary } from './cloud-connection.js'; + +export type QuotaLibrary = ReturnType; + +export const createQuotaLibrary = (queries: Queries, cloudConnection: CloudConnectionLibrary) => { + const { + applications: { countNonM2mApplications, countM2mApplications }, + resources: { findTotalNumberOfResources }, + } = queries; + + const getTenantUsage = async (key: keyof FeatureQuota): Promise => { + if (key === 'applicationsLimit') { + return countNonM2mApplications(); + } + + if (key === 'machineToMachineLimit') { + return countM2mApplications(); + } + + if (key === 'resourcesLimit') { + const { count } = await findTotalNumberOfResources(); + // Ignore the default management API resource + return count - 1; + } + + // TODO: add other keys + + throw new Error('Unsupported subscription quota key'); + }; + + const guardKey = async (key: keyof FeatureQuota) => { + const { isCloud, isIntegrationTest, isProduction } = EnvSet.values; + + // Cloud only feature, skip in non-cloud production environments + if (isProduction && !isCloud) { + return; + } + + // Disable in integration tests + if (isIntegrationTest) { + return; + } + + // TODO @sijie: remove this when pricing is ready + if (isProduction) { + return; + } + + const plan = await getTenantSubscriptionPlan(cloudConnection); + const limit = plan.quota[key]; + + if (typeof limit === 'boolean') { + assertThat( + limit, + new RequestError({ + code: 'subscription.limit_exceeded', + status: 403, + data: { + key, + }, + }) + ); + } else if (typeof limit === 'number') { + const tenantUsage = await getTenantUsage(key); + + assertThat( + tenantUsage < limit, + new RequestError({ + code: 'subscription.limit_exceeded', + status: 403, + data: { + key, + limit, + usage: tenantUsage, + }, + }) + ); + } else { + throw new TypeError('Unsupported subscription quota type'); + } + }; + + return { guardKey }; +}; diff --git a/packages/core/src/middleware/koa-quota-guard.test.ts b/packages/core/src/middleware/koa-quota-guard.test.ts deleted file mode 100644 index a6a3eed21..000000000 --- a/packages/core/src/middleware/koa-quota-guard.test.ts +++ /dev/null @@ -1,97 +0,0 @@ -import { GlobalValues } from '@logto/shared'; -import { createMockUtils } from '@logto/shared/esm'; -import { type Context } from 'koa'; - -import { mockFreePlan } from '#src/__mocks__/subscription.js'; -import { createMockCloudConnectionLibrary } from '#src/test-utils/cloud-connection.js'; -import createMockContext from '#src/test-utils/jest-koa-mocks/create-mock-context.js'; -import { MockQueries } from '#src/test-utils/tenant.js'; - -const { jest } = import.meta; - -const { mockEsmWithActual } = createMockUtils(jest); - -const getValues = jest.fn(() => ({ - ...new GlobalValues(), - isCloud: true, -})); - -await mockEsmWithActual('#src/env-set/index.js', () => ({ - EnvSet: { - get values() { - return getValues(); - }, - }, -})); - -const { getTenantSubscriptionPlan } = await mockEsmWithActual( - '#src/utils/subscription/index.js', - () => ({ - getTenantSubscriptionPlan: jest.fn().mockResolvedValue(mockFreePlan), - }) -); - -const { default: koaQuotaGuard } = await import('./koa-quota-guard.js'); - -const createContext = (): Context => { - return createMockContext(); -}; - -const countNonM2MApplications = jest.fn(); -const queries = new MockQueries({ - applications: { countNonM2MApplications }, -}); - -const cloudConnection = createMockCloudConnectionLibrary(); - -describe('koaQuotaGuard() middleware', () => { - afterEach(() => { - getTenantSubscriptionPlan.mockClear(); - getValues.mockReturnValue({ - ...new GlobalValues(), - isCloud: true, - }); - }); - - it('should skip on non-cloud', async () => { - getValues.mockReturnValueOnce({ - ...new GlobalValues(), - isCloud: false, - }); - - const ctx = createContext(); - await koaQuotaGuard({ - key: 'applicationsLimit', - queries, - cloudConnection, - })(ctx, jest.fn()); - - expect(getTenantSubscriptionPlan).not.toHaveBeenCalled(); - }); - - it('should pass when limit is not exeeded', async () => { - countNonM2MApplications.mockResolvedValueOnce(0); - - const ctx = createContext(); - await expect( - koaQuotaGuard({ - key: 'applicationsLimit', - queries, - cloudConnection, - })(ctx, jest.fn()) - ).resolves.not.toThrow(); - }); - - it('should throw when limit is exeeded', async () => { - countNonM2MApplications.mockResolvedValueOnce(mockFreePlan.quota.applicationsLimit); - - const ctx = createContext(); - await expect( - koaQuotaGuard({ - key: 'applicationsLimit', - queries, - cloudConnection, - })(ctx, jest.fn()) - ).rejects.toThrow(); - }); -}); diff --git a/packages/core/src/middleware/koa-quota-guard.ts b/packages/core/src/middleware/koa-quota-guard.ts index 6a74ce23b..dd6836aff 100644 --- a/packages/core/src/middleware/koa-quota-guard.ts +++ b/packages/core/src/middleware/koa-quota-guard.ts @@ -1,75 +1,19 @@ import type { MiddlewareType } from 'koa'; -import { EnvSet } from '#src/env-set/index.js'; -import RequestError from '#src/errors/RequestError/index.js'; -import { type CloudConnectionLibrary } from '#src/libraries/cloud-connection.js'; -import type Queries from '#src/tenants/Queries.js'; -import assertThat from '#src/utils/assert-that.js'; -import { getTenantSubscriptionPlan } from '#src/utils/subscription/index.js'; +import { type QuotaLibrary } from '#src/libraries/quota.js'; import { type FeatureQuota } from '#src/utils/subscription/types.js'; type UsageGuardConfig = { key: keyof FeatureQuota; - cloudConnection: CloudConnectionLibrary; - queries: Queries; -}; - -const getTenantUsage = async (key: keyof FeatureQuota, queries: Queries): Promise => { - if (key === 'applicationsLimit') { - return queries.applications.countNonM2MApplications(); - } - - // TODO: add other keys - - throw new Error('Unsupported subscription quota key'); + quota: QuotaLibrary; }; export default function koaQuotaGuard({ key, - queries, - cloudConnection, + quota, }: UsageGuardConfig): MiddlewareType { return async (ctx, next) => { - const { isCloud, isIntegrationTest, isProduction } = EnvSet.values; - - // Disable in production until pricing is ready - if (!isCloud || isIntegrationTest || isProduction) { - return next(); - } - - const plan = await getTenantSubscriptionPlan(cloudConnection); - const limit = plan.quota[key]; - - if (typeof limit === 'boolean') { - assertThat( - limit, - new RequestError({ - code: 'subscription.limit_exceeded', - status: 403, - data: { - key, - }, - }) - ); - } else if (typeof limit === 'number') { - const tenantUsage = await getTenantUsage(key, queries); - - assertThat( - tenantUsage < limit, - new RequestError({ - code: 'subscription.limit_exceeded', - status: 403, - data: { - key, - limit, - usage: tenantUsage, - }, - }) - ); - } else { - throw new TypeError('Unsupported subscription quota type'); - } - + await quota.guardKey(key); return next(); }; } diff --git a/packages/core/src/queries/application.ts b/packages/core/src/queries/application.ts index c8d5b9f89..c03b85a64 100644 --- a/packages/core/src/queries/application.ts +++ b/packages/core/src/queries/application.ts @@ -28,7 +28,7 @@ export const createApplicationQueries = (pool: CommonQueryMethods) => { id: string, set: Partial> ) => updateApplication({ set, where: { id }, jsonbMode: 'merge' }); - const countNonM2MApplications = async () => { + const countNonM2mApplications = async () => { const { count } = await pool.one<{ count: string }>(sql` select count(*) from ${table} @@ -37,6 +37,15 @@ export const createApplicationQueries = (pool: CommonQueryMethods) => { return Number(count); }; + const countM2mApplications = async () => { + const { count } = await pool.one<{ count: string }>(sql` + select count(*) + from ${table} + where ${fields.type} = ${ApplicationType.MachineToMachine} + `); + + return Number(count); + }; const deleteApplicationById = async (id: string) => { const { rowCount } = await pool.query(sql` @@ -56,7 +65,8 @@ export const createApplicationQueries = (pool: CommonQueryMethods) => { insertApplication, updateApplication, updateApplicationById, - countNonM2MApplications, + countNonM2mApplications, + countM2mApplications, deleteApplicationById, }; }; diff --git a/packages/core/src/routes/application.test.ts b/packages/core/src/routes/application.test.ts index 7a73dd3f3..b61d99f2e 100644 --- a/packages/core/src/routes/application.test.ts +++ b/packages/core/src/routes/application.test.ts @@ -3,6 +3,7 @@ import { ApplicationType } from '@logto/schemas'; import { pickDefault, createMockUtils } from '@logto/shared/esm'; import { mockApplication } from '#src/__mocks__/index.js'; +import { createMockQuotaLibrary } from '#src/test-utils/quota.js'; import { MockTenant } from '#src/test-utils/tenant.js'; const { jest } = import.meta; @@ -17,30 +18,35 @@ await mockEsmWithActual('@logto/shared', () => ({ generateStandardId: () => 'randomId', })); -const tenantContext = new MockTenant(undefined, { - applications: { - findTotalNumberOfApplications: jest.fn(async () => ({ count: 10 })), - findAllApplications: jest.fn(async () => [mockApplication]), - findApplicationById, - deleteApplicationById, - insertApplication: jest.fn( - async (body: CreateApplication): Promise => ({ - ...mockApplication, - ...body, - oidcClientMetadata: { - ...mockApplication.oidcClientMetadata, - ...body.oidcClientMetadata, - }, - }) - ), - updateApplicationById: jest.fn( - async (_, data: Partial): Promise => ({ - ...mockApplication, - ...data, - }) - ), +const tenantContext = new MockTenant( + undefined, + { + applications: { + findTotalNumberOfApplications: jest.fn(async () => ({ count: 10 })), + findAllApplications: jest.fn(async () => [mockApplication]), + findApplicationById, + deleteApplicationById, + insertApplication: jest.fn( + async (body: CreateApplication): Promise => ({ + ...mockApplication, + ...body, + oidcClientMetadata: { + ...mockApplication.oidcClientMetadata, + ...body.oidcClientMetadata, + }, + }) + ), + updateApplicationById: jest.fn( + async (_, data: Partial): Promise => ({ + ...mockApplication, + ...data, + }) + ), + }, }, -}); + undefined, + { quota: createMockQuotaLibrary() } +); const { createRequester } = await import('#src/utils/test-utils.js'); const applicationRoutes = await pickDefault(import('./application.js')); diff --git a/packages/core/src/routes/application.ts b/packages/core/src/routes/application.ts index f096718e6..e5cf2d89c 100644 --- a/packages/core/src/routes/application.ts +++ b/packages/core/src/routes/application.ts @@ -4,6 +4,7 @@ import { buildDemoAppDataForTenant, Applications, InternalRole, + ApplicationType, } from '@logto/schemas'; import { generateStandardId, buildIdGenerator } from '@logto/shared'; import { boolean, object, string, z } from 'zod'; @@ -11,7 +12,6 @@ import { boolean, object, string, z } from 'zod'; import RequestError from '#src/errors/RequestError/index.js'; import koaGuard from '#src/middleware/koa-guard.js'; import koaPagination from '#src/middleware/koa-pagination.js'; -import koaQuotaGuard from '#src/middleware/koa-quota-guard.js'; import { buildOidcClientMetadata } from '#src/oidc/utils.js'; import assertThat from '#src/utils/assert-that.js'; @@ -22,7 +22,14 @@ const includesInternalAdminRole = (roles: Readonly>) => roles.some(({ role: { name } }) => name === InternalRole.Admin); export default function applicationRoutes( - ...[router, { queries, id: tenantId, cloudConnection }]: RouterInitArgs + ...[ + router, + { + queries, + id: tenantId, + libraries: { quota }, + }, + ]: RouterInitArgs ) { const { deleteApplicationById, @@ -64,7 +71,6 @@ export default function applicationRoutes( router.post( '/applications', - koaQuotaGuard({ key: 'applicationsLimit', cloudConnection, queries }), koaGuard({ body: Applications.createGuard .omit({ id: true, createdAt: true }) @@ -76,6 +82,12 @@ export default function applicationRoutes( async (ctx, next) => { const { oidcClientMetadata, ...rest } = ctx.guard.body; + await quota.guardKey( + rest.type === ApplicationType.MachineToMachine + ? 'machineToMachineLimit' + : 'applicationsLimit' + ); + ctx.body = await insertApplication({ id: applicationId(), secret: generateStandardId(), diff --git a/packages/core/src/routes/resource.test.ts b/packages/core/src/routes/resource.test.ts index ad9cc54cf..e790a7a59 100644 --- a/packages/core/src/routes/resource.test.ts +++ b/packages/core/src/routes/resource.test.ts @@ -3,6 +3,7 @@ import { pickDefault, createMockUtils } from '@logto/shared/esm'; import { type Nullable } from '@silverhand/essentials'; import { mockResource, mockScope } from '#src/__mocks__/index.js'; +import { createMockQuotaLibrary } from '#src/test-utils/quota.js'; import { MockTenant } from '#src/test-utils/tenant.js'; import { createRequester } from '#src/utils/test-utils.js'; @@ -52,6 +53,7 @@ const libraries = { scopes: [], })), }, + quota: createMockQuotaLibrary(), }; mockEsm('@logto/shared', () => ({ diff --git a/packages/core/src/routes/resource.ts b/packages/core/src/routes/resource.ts index 6f487e71a..378728725 100644 --- a/packages/core/src/routes/resource.ts +++ b/packages/core/src/routes/resource.ts @@ -6,6 +6,7 @@ import { boolean, object, string } from 'zod'; import RequestError from '#src/errors/RequestError/index.js'; import koaGuard from '#src/middleware/koa-guard.js'; import koaPagination from '#src/middleware/koa-pagination.js'; +import koaQuotaGuard from '#src/middleware/koa-quota-guard.js'; import assertThat from '#src/utils/assert-that.js'; import { parseSearchParamsForSearch } from '#src/utils/search.js'; @@ -76,6 +77,7 @@ export default function resourceRoutes( router.post( '/resources', + koaQuotaGuard({ key: 'resourcesLimit', quota: libraries.quota }), koaGuard({ // Intentionally omit `isDefault` since it'll affect other rows. // Use the dedicated API `PATCH /resources/:id/is-default` to update. diff --git a/packages/core/src/tenants/Libraries.ts b/packages/core/src/tenants/Libraries.ts index 7a608b5e2..639436357 100644 --- a/packages/core/src/tenants/Libraries.ts +++ b/packages/core/src/tenants/Libraries.ts @@ -1,9 +1,11 @@ import { createApplicationLibrary } from '#src/libraries/application.js'; +import { type CloudConnectionLibrary } from '#src/libraries/cloud-connection.js'; import type { ConnectorLibrary } from '#src/libraries/connector.js'; import { createDomainLibrary } from '#src/libraries/domain.js'; import { createHookLibrary } from '#src/libraries/hook/index.js'; import { createPasscodeLibrary } from '#src/libraries/passcode.js'; import { createPhraseLibrary } from '#src/libraries/phrase.js'; +import { createQuotaLibrary } from '#src/libraries/quota.js'; import { createResourceLibrary } from '#src/libraries/resource.js'; import { createSignInExperienceLibrary } from '#src/libraries/sign-in-experience/index.js'; import { createSocialLibrary } from '#src/libraries/social.js'; @@ -23,11 +25,13 @@ export default class Libraries { applications = createApplicationLibrary(this.queries); verificationStatuses = createVerificationStatusLibrary(this.queries); domains = createDomainLibrary(this.queries); + quota = createQuotaLibrary(this.queries, this.cloudConnection); constructor( public readonly tenantId: string, private readonly queries: Queries, // Explicitly passing connector library to eliminate dependency issue - private readonly connectors: ConnectorLibrary + private readonly connectors: ConnectorLibrary, + private readonly cloudConnection: CloudConnectionLibrary ) {} } diff --git a/packages/core/src/tenants/Tenant.ts b/packages/core/src/tenants/Tenant.ts index 1b5609308..8e51819a7 100644 --- a/packages/core/src/tenants/Tenant.ts +++ b/packages/core/src/tenants/Tenant.ts @@ -57,7 +57,7 @@ export default class Tenant implements TenantContext { public readonly logtoConfigs = createLogtoConfigLibrary(queries), public readonly cloudConnection = createCloudConnectionLibrary(logtoConfigs), public readonly connectors = createConnectorLibrary(queries, cloudConnection), - public readonly libraries = new Libraries(id, queries, connectors) + public readonly libraries = new Libraries(id, queries, connectors, cloudConnection) ) { const isAdminTenant = id === adminTenantId; const mountedApps = [ diff --git a/packages/core/src/test-utils/quota.ts b/packages/core/src/test-utils/quota.ts new file mode 100644 index 000000000..3dbb88b80 --- /dev/null +++ b/packages/core/src/test-utils/quota.ts @@ -0,0 +1,9 @@ +import { type QuotaLibrary } from '#src/libraries/quota.js'; + +const { jest } = import.meta; + +export const createMockQuotaLibrary = (): QuotaLibrary => { + return { + guardKey: jest.fn(), + }; +}; diff --git a/packages/core/src/test-utils/tenant.ts b/packages/core/src/test-utils/tenant.ts index 002f547cd..3d0be39e8 100644 --- a/packages/core/src/test-utils/tenant.ts +++ b/packages/core/src/test-utils/tenant.ts @@ -79,7 +79,7 @@ export class MockTenant implements TenantContext { ...createConnectorLibrary(this.queries, this.cloudConnection), ...connectorsOverride, }; - this.libraries = new Libraries(this.id, this.queries, this.connectors); + this.libraries = new Libraries(this.id, this.queries, this.connectors, this.cloudConnection); this.setPartial('libraries', librariesOverride); }