Infrastructure/logto
0
Fork 0
mirror of https://github.com/logto-io/logto.git synced 2025-01-27 21:39:16 -05:00
Code Issues Activity

refactor(core): create quota library (#4185)

Browse source
This commit is contained in:
wangsijie 2023-07-20 16:29:39 +08:00 • committed by GitHub
parent 7cda535629
commit d5885160cc
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
14 changed files with 222 additions and 190 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
packages/core/src/database/row-count.ts
View file

@ -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) };
};

44
packages/core/src/libraries/quota.test.ts Normal file
View file

@ -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();
});
});

91
packages/core/src/libraries/quota.ts Normal file
View file

@ -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<typeof createQuotaLibrary>;
export const createQuotaLibrary = (queries: Queries, cloudConnection: CloudConnectionLibrary) => {
const {
applications: { countNonM2mApplications, countM2mApplications },
resources: { findTotalNumberOfResources },
} = queries;
const getTenantUsage = async (key: keyof FeatureQuota): Promise<number> => {
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 };
};

97
packages/core/src/middleware/koa-quota-guard.test.ts
View file

@ -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();
});
});

64
packages/core/src/middleware/koa-quota-guard.ts
View file

@ -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<number> => {
if (key === 'applicationsLimit') {
return queries.applications.countNonM2MApplications();
}
// TODO: add other keys
throw new Error('Unsupported subscription quota key');
quota: QuotaLibrary;
};
export default function koaQuotaGuard<StateT, ContextT, ResponseBodyT>({
key,
queries,
cloudConnection,
quota,
}: UsageGuardConfig): MiddlewareType<StateT, ContextT, ResponseBodyT> {
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();
};
}

14
packages/core/src/queries/application.ts
View file

@ -28,7 +28,7 @@ export const createApplicationQueries = (pool: CommonQueryMethods) => {
id: string,
set: Partial<OmitAutoSetFields<CreateApplication>>
) => 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,
};
};

52
packages/core/src/routes/application.test.ts
View file

@ -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<Application> => ({
...mockApplication,
...body,
oidcClientMetadata: {
...mockApplication.oidcClientMetadata,
...body.oidcClientMetadata,
},
})
),
updateApplicationById: jest.fn(
async (_, data: Partial<CreateApplication>): Promise<Application> => ({
...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<Application> => ({
...mockApplication,
...body,
oidcClientMetadata: {
...mockApplication.oidcClientMetadata,
...body.oidcClientMetadata,
},
})
),
updateApplicationById: jest.fn(
async (_, data: Partial<CreateApplication>): Promise<Application> => ({
...mockApplication,
...data,
})
),
},
},
});
undefined,
{ quota: createMockQuotaLibrary() }
);
const { createRequester } = await import('#src/utils/test-utils.js');
const applicationRoutes = await pickDefault(import('./application.js'));

18
packages/core/src/routes/application.ts
View file

@ -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<Array<{ role: Role }>>) =>
roles.some(({ role: { name } }) => name === InternalRole.Admin);
export default function applicationRoutes<T extends AuthedRouter>(
...[router, { queries, id: tenantId, cloudConnection }]: RouterInitArgs<T>
...[
router,
{
queries,
id: tenantId,
libraries: { quota },
},
]: RouterInitArgs<T>
) {
const {
deleteApplicationById,
@ -64,7 +71,6 @@ export default function applicationRoutes<T extends AuthedRouter>(
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<T extends AuthedRouter>(
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(),

2
packages/core/src/routes/resource.test.ts
View file

@ -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', () => ({

2
packages/core/src/routes/resource.ts
View file

@ -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<T extends AuthedRouter>(
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.

6
packages/core/src/tenants/Libraries.ts
View file

@ -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
) {}
}

2
packages/core/src/tenants/Tenant.ts
View file

@ -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 = [

9
packages/core/src/test-utils/quota.ts Normal file
View file

@ -0,0 +1,9 @@
import { type QuotaLibrary } from '#src/libraries/quota.js';
const { jest } = import.meta;
export const createMockQuotaLibrary = (): QuotaLibrary => {
return {
guardKey: jest.fn(),
};
};

2
packages/core/src/test-utils/tenant.ts
View file

@ -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);
}