Significantly faster users/[user_id] endpoint (and some others) (#998)

This commit is contained in:
Konsti Wohlwend 2025-11-05 09:15:36 -08:00 committed by GitHub
parent e824e8c513
commit fbf36d1004
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 235 additions and 60 deletions

View File

@ -1,8 +1,8 @@
import { getRenderedEnvironmentConfigQuery } from "@/lib/config";
import { getRenderedProjectConfigQuery } from "@/lib/config";
import { normalizeEmail } from "@/lib/emails";
import { grantDefaultProjectPermissions } from "@/lib/permissions";
import { ensureTeamMembershipExists, ensureUserExists } from "@/lib/request-checks";
import { Tenancy, getSoleTenancyFromProjectBranch, getTenancy } from "@/lib/tenancies";
import { Tenancy, getSoleTenancyFromProjectBranch } from "@/lib/tenancies";
import { PrismaTransaction } from "@/lib/types";
import { sendTeamMembershipDeletedWebhook, sendUserCreatedWebhook, sendUserDeletedWebhook, sendUserUpdatedWebhook } from "@/lib/webhooks";
import { RawQuery, getPrismaClientForSourceOfTruth, getPrismaClientForTenancy, getPrismaSchemaForSourceOfTruth, getPrismaSchemaForTenancy, globalPrismaClient, rawQuery, retryTransaction, sqlQuoteIdent } from "@/prisma-client";
@ -388,20 +388,21 @@ export function getUserIfOnGlobalPrismaClientQuery(projectId: string, branchId:
};
}
export async function getUser(options: { userId: string } & ({ projectId: string, branchId: string } | { tenancyId: string })) {
let projectId, branchId;
if (!("tenancyId" in options)) {
export async function getUser(options: { userId: string } & ({ projectId: string, branchId: string } | { tenancy: Tenancy })) {
let projectId, branchId, sourceOfTruth;
if ("tenancy" in options) {
projectId = options.tenancy.project.id;
branchId = options.tenancy.branchId;
sourceOfTruth = options.tenancy.config.sourceOfTruth;
} else {
projectId = options.projectId;
branchId = options.branchId;
} else {
const tenancy = await getTenancy(options.tenancyId) ?? throwErr("Tenancy not found", { tenancyId: options.tenancyId });
projectId = tenancy.project.id;
branchId = tenancy.branchId;
const projectConfig = await rawQuery(globalPrismaClient, getRenderedProjectConfigQuery({ projectId }));
sourceOfTruth = projectConfig.sourceOfTruth;
}
const environmentConfig = await rawQuery(globalPrismaClient, getRenderedEnvironmentConfigQuery({ projectId, branchId }));
const prisma = await getPrismaClientForSourceOfTruth(environmentConfig.sourceOfTruth, branchId);
const schema = await getPrismaSchemaForSourceOfTruth(environmentConfig.sourceOfTruth, branchId);
const prisma = await getPrismaClientForSourceOfTruth(sourceOfTruth, branchId);
const schema = await getPrismaSchemaForSourceOfTruth(sourceOfTruth, branchId);
const result = await rawQuery(prisma, getUserQuery(projectId, branchId, options.userId, schema));
return result;
}
@ -421,7 +422,7 @@ export const usersCrudHandlers = createLazyProxy(() => createCrudHandlers(usersC
include_anonymous: yupString().oneOf(["true", "false"]).optional().meta({ openapiField: { onlyShowInOperations: [ 'List' ], description: "Whether to include anonymous users in the results. Defaults to false" } }),
}),
onRead: async ({ auth, params, query }) => {
const user = await getUser({ tenancyId: auth.tenancy.id, userId: params.user_id });
const user = await getUser({ tenancy: auth.tenancy, userId: params.user_id });
if (!user) {
throw new KnownErrors.UserNotFound();
}

View File

@ -131,10 +131,7 @@ export function getProjectConfigOverrideQuery(options: ProjectOptions): RawQuery
if (queryResult.length > 1) {
throw new StackAssertionError(`Expected 0 or 1 project config overrides for project ${options.projectId}, got ${queryResult.length}`, { queryResult });
}
if (queryResult.length === 0) {
throw new StackAssertionError(`Expected a project row for project ${options.projectId}, got 0`, { queryResult, options });
}
return migrateConfigOverride("project", queryResult[0].projectConfigOverride ?? {});
return migrateConfigOverride("project", queryResult[0]?.projectConfigOverride ?? {});
},
};
}
@ -142,13 +139,13 @@ export function getProjectConfigOverrideQuery(options: ProjectOptions): RawQuery
export function getBranchConfigOverrideQuery(options: BranchOptions): RawQuery<Promise<BranchConfigOverride>> {
// fetch branch config from GitHub
// (currently it's just empty)
if (options.branchId !== DEFAULT_BRANCH_ID) {
throw new StackAssertionError('Not implemented');
}
return {
supportedPrismaClients: ["global"],
sql: Prisma.sql`SELECT 1`,
postProcess: async () => {
if (options.branchId !== DEFAULT_BRANCH_ID) {
throw new StackAssertionError('getBranchConfigOverrideQuery is not implemented for branches other than the default one');
}
return migrateConfigOverride("branch", {});
},
};

View File

@ -1,9 +1,11 @@
import { globalPrismaClient, rawQuery } from "@/prisma-client";
import { globalPrismaClient, RawQuery, rawQuery } from "@/prisma-client";
import { Prisma } from "@prisma/client";
import { ProjectsCrud } from "@stackframe/stack-shared/dist/interface/crud/projects";
import { getNodeEnvironment } from "@stackframe/stack-shared/dist/utils/env";
import { StackAssertionError, throwErr } from "@stackframe/stack-shared/dist/utils/errors";
import { deepPlainEquals } from "@stackframe/stack-shared/dist/utils/objects";
import { getRenderedOrganizationConfigQuery } from "./config";
import { getProject } from "./projects";
import { getProject, getProjectQuery } from "./projects";
/**
* @deprecated YOU PROBABLY ALMOST NEVER WANT TO USE THIS, UNLESS YOU ACTUALLY NEED THE DEFAULT BRANCH ID. DON'T JUST USE THIS TO GET A TENANCY BECAUSE YOU DON'T HAVE ONE
@ -16,7 +18,11 @@ import { getProject } from "./projects";
*/
export const DEFAULT_BRANCH_ID = "main";
export async function tenancyPrismaToCrud(prisma: Prisma.TenancyGetPayload<{}>) {
/**
* @deprecated UNUSED: This function is only kept for development mode validation in getTenancyFromProject.
* The old Prisma-based implementation, replaced by getTenancyFromProjectQuery which uses RawQuery.
*/
async function tenancyPrismaToCrudUnused(prisma: Prisma.TenancyGetPayload<{}>) {
if (prisma.hasNoOrganization && prisma.organizationId !== null) {
throw new StackAssertionError("Organization ID is not null for a tenancy with hasNoOrganization", { tenancyId: prisma.id, prisma });
}
@ -44,7 +50,7 @@ export async function tenancyPrismaToCrud(prisma: Prisma.TenancyGetPayload<{}>)
};
}
export type Tenancy = Awaited<ReturnType<typeof tenancyPrismaToCrud>>;
export type Tenancy = Awaited<ReturnType<typeof tenancyPrismaToCrudUnused>>;
/**
* @deprecated This is a temporary function for the situation where every project-branch has exactly one tenancy. Later,
@ -57,7 +63,7 @@ export function getSoleTenancyFromProjectBranch(project: Omit<ProjectsCrud["Admi
*/
export function getSoleTenancyFromProjectBranch(project: Omit<ProjectsCrud["Admin"]["Read"], "config"> | string, branchId: string, returnNullIfNotFound: boolean): Promise<Tenancy | null>;
export async function getSoleTenancyFromProjectBranch(projectOrId: Omit<ProjectsCrud["Admin"]["Read"], "config"> | string, branchId: string, returnNullIfNotFound: boolean = false): Promise<Tenancy | null> {
const res = await getTenancyFromProject(typeof projectOrId === 'string' ? projectOrId : projectOrId.id, branchId, null);
const res = await rawQuery(globalPrismaClient, getSoleTenancyFromProjectBranchQuery(projectOrId, branchId, true));
if (!res) {
if (returnNullIfNotFound) return null;
throw new StackAssertionError(`No tenancy found for project ${typeof projectOrId === 'string' ? projectOrId : projectOrId.id}`, { projectOrId });
@ -65,6 +71,14 @@ export async function getSoleTenancyFromProjectBranch(projectOrId: Omit<Projects
return res;
}
/**
* @deprecated This is a temporary function for the situation where every project-branch has exactly one tenancy. Later,
* we will support multiple tenancies per project-branch, and all uses of this function will be refactored.
*/
export function getSoleTenancyFromProjectBranchQuery(project: Omit<ProjectsCrud["Admin"]["Read"], "config"> | string, branchId: string, returnNullIfNotFound: true): RawQuery<Promise<Tenancy | null>> {
return getTenancyFromProjectQuery(typeof project === 'string' ? project : project.id, branchId, null);
}
export async function getTenancy(tenancyId: string) {
if (tenancyId === "internal") {
throw new StackAssertionError("Tried to get tenancy with ID `internal`. This is a mistake because `internal` is only a valid identifier for projects.");
@ -73,31 +87,135 @@ export async function getTenancy(tenancyId: string) {
where: { id: tenancyId },
});
if (!prisma) return null;
return await tenancyPrismaToCrud(prisma);
return await getTenancyFromProject(prisma.projectId, prisma.branchId, prisma.organizationId);
}
function getTenancyFromProjectQuery(projectId: string, branchId: string, organizationId: string | null): RawQuery<Promise<Tenancy | null>> {
return RawQuery.then(
RawQuery.all([
{
supportedPrismaClients: ["global"],
sql: organizationId === null
? Prisma.sql`
SELECT "Tenancy".*
FROM "Tenancy"
WHERE "Tenancy"."projectId" = ${projectId}
AND "Tenancy"."branchId" = ${branchId}
AND "Tenancy"."hasNoOrganization" = 'TRUE'
`
: Prisma.sql`
SELECT "Tenancy".*
FROM "Tenancy"
WHERE "Tenancy"."projectId" = ${projectId}
AND "Tenancy"."branchId" = ${branchId}
AND "Tenancy"."organizationId" = ${organizationId}
`,
postProcess: (queryResult) => {
if (queryResult.length > 1) {
throw new StackAssertionError(
`Expected 0 or 1 tenancies for project ${projectId}, branch ${branchId}, organization ${organizationId}, got ${queryResult.length}`,
{ queryResult }
);
}
if (queryResult.length === 0) {
return Promise.resolve(null);
}
return Promise.resolve(queryResult[0] as Prisma.TenancyGetPayload<{}>);
},
},
getProjectQuery(projectId),
getRenderedOrganizationConfigQuery({
projectId,
branchId,
organizationId,
}),
] as const),
async ([tenancyResultPromise, projectResultPromise, configPromise]) => {
const tenancyResult = await tenancyResultPromise;
if (!tenancyResult) return null;
const [projectResult, config] = await Promise.all([
projectResultPromise,
configPromise,
]);
if (!projectResult) {
throw new StackAssertionError("Project in tenancy not found", { projectId, tenancyId: tenancyResult.id });
}
// Validate tenancy consistency
if (tenancyResult.hasNoOrganization && tenancyResult.organizationId !== null) {
throw new StackAssertionError("Organization ID is not null for a tenancy with hasNoOrganization", {
tenancyId: tenancyResult.id,
tenancy: tenancyResult
});
}
if (!tenancyResult.hasNoOrganization && tenancyResult.organizationId === null) {
throw new StackAssertionError("Organization ID is null for a tenancy without hasNoOrganization", {
tenancyId: tenancyResult.id,
tenancy: tenancyResult
});
}
return {
id: tenancyResult.id,
config,
branchId: tenancyResult.branchId,
organization: tenancyResult.organizationId === null ? null : {
// TODO actual organization type
id: tenancyResult.organizationId,
},
project: projectResult,
};
}
);
}
/**
* @deprecated Not actually deprecated but if you're using this you're probably doing something wrong ask Konsti for help
*
* (if Konsti is not around unless you are editing the implementation of SmartRequestAuth, you should probably take the
* tenancy from the SmartRequest auth parameter instead of fetching your own. If you are editing the SmartRequestAuth
* implementation carry on.)
*/
export async function getTenancyFromProject(projectId: string, branchId: string, organizationId: string | null) {
const prisma = await globalPrismaClient.tenancy.findUnique({
where: {
...(organizationId === null ? {
projectId_branchId_hasNoOrganization: {
projectId: projectId,
branchId: branchId,
hasNoOrganization: "TRUE",
}
} : {
projectId_branchId_organizationId: {
projectId: projectId,
branchId: branchId,
organizationId: organizationId,
}
}),
},
});
if (!prisma) return null;
return await tenancyPrismaToCrud(prisma);
// Use the new RawQuery implementation
const result = await rawQuery(globalPrismaClient, getTenancyFromProjectQuery(projectId, branchId, organizationId));
// In development mode, compare with the old implementation to ensure correctness
if (!getNodeEnvironment().includes("prod")) {
const prisma = await globalPrismaClient.tenancy.findUnique({
where: {
...(organizationId === null ? {
projectId_branchId_hasNoOrganization: {
projectId: projectId,
branchId: branchId,
hasNoOrganization: "TRUE",
}
} : {
projectId_branchId_organizationId: {
projectId: projectId,
branchId: branchId,
organizationId: organizationId,
}
}),
},
});
const oldResult = prisma ? await tenancyPrismaToCrudUnused(prisma) : null;
// Compare the two results
if (!deepPlainEquals(result, oldResult)) {
throw new StackAssertionError("getTenancyFromProject: new implementation does not match old implementation", {
projectId,
branchId,
organizationId,
newResult: result,
oldResult,
});
}
}
return result;
}

View File

@ -4,7 +4,7 @@ import { getUser, getUserIfOnGlobalPrismaClientQuery } from "@/app/api/latest/us
import { getRenderedEnvironmentConfigQuery } from "@/lib/config";
import { checkApiKeySet, checkApiKeySetQuery } from "@/lib/internal-api-keys";
import { getProjectQuery, listManagedProjectIds } from "@/lib/projects";
import { DEFAULT_BRANCH_ID, Tenancy, getSoleTenancyFromProjectBranch } from "@/lib/tenancies";
import { DEFAULT_BRANCH_ID, Tenancy, getSoleTenancyFromProjectBranchQuery } from "@/lib/tenancies";
import { decodeAccessToken } from "@/lib/tokens";
import { globalPrismaClient, rawQueryAll } from "@/prisma-client";
import { KnownErrors } from "@stackframe/stack-shared";
@ -248,23 +248,13 @@ const parseAuth = withTraceSpan('smart request parseAuth', async (req: NextReque
isServerKeyValid: secretServerKey && requestType === "server" ? checkApiKeySetQuery(projectId, { secretServerKey }) : undefined,
isAdminKeyValid: superSecretAdminKey && requestType === "admin" ? checkApiKeySetQuery(projectId, { superSecretAdminKey }) : undefined,
project: getProjectQuery(projectId),
tenancy: getSoleTenancyFromProjectBranchQuery(projectId, branchId, true),
environmentRenderedConfig: getRenderedEnvironmentConfigQuery({ projectId, branchId }),
};
const queriesResults = await rawQueryAll(globalPrismaClient, bundledQueries);
const project = await queriesResults.project;
if (project === null) throw new KnownErrors.CurrentProjectNotFound(projectId); // this does allow one to probe whether a project exists or not, but that's fine
const environmentConfig = await queriesResults.environmentRenderedConfig;
// As explained above, as a performance optimization we already fetch the user from the global database optimistically
// If it turned out that the source-of-truth is not the global database, we'll fetch the user from the source-of-truth
// database instead.
const user = environmentConfig.sourceOfTruth.type === "hosted"
? queriesResults.userIfOnGlobalPrismaClient
: (userId ? await getUser({ userId, projectId, branchId }) : undefined);
// TODO HACK tenancy is not needed for /users/me, so let's not fetch it as a hack to make the endpoint faster. Once we
// refactor this stuff, we can fetch the tenancy alongside the user and won't need this anymore
const tenancy = req.method === "GET" && req.url.endsWith("/users/me") ? "tenancy not available in /users/me as a performance hack" as never : await getSoleTenancyFromProjectBranch(projectId, branchId, true);
if (project === null) throw new KnownErrors.CurrentProjectNotFound(projectId); // this does allow one to probe whether a project exists or not, but that's fine (it's worth the better error messages)
const tenancy = await queriesResults.tenancy;
if (developmentKeyOverride) {
if (!["development", "test"].includes(getNodeEnvironment()) && getEnvVariable("STACK_ALLOW_DEVELOPMENT_KEY_OVERRIDE_DESPITE_PRODUCTION", "") !== "this-is-dangerous") { // it's not actually that dangerous, but it changes the security model
@ -299,9 +289,17 @@ const parseAuth = withTraceSpan('smart request parseAuth', async (req: NextReque
}
if (!tenancy) {
// note that we only check branch existence here so you can't probe branches unless you have the project keys
throw new KnownErrors.BranchDoesNotExist(branchId);
}
// As explained above, as a performance optimization we already fetch the user from the global database optimistically
// If it turned out that the source-of-truth is not the global database, we'll fetch the user from the source-of-truth
// database instead.
const user = tenancy.config.sourceOfTruth.type === "hosted"
? queriesResults.userIfOnGlobalPrismaClient
: (userId ? await getUser({ userId, projectId, branchId }) : undefined);
return {
project,
branchId,

View File

@ -67,6 +67,7 @@ export function createMailbox(email?: string): Mailbox {
export type ProjectKeys = "no-project" | {
projectId: string,
branchId?: string,
publishableClientKey?: string,
secretServerKey?: string,
superSecretAdminKey?: string,

View File

@ -76,6 +76,66 @@ describe("without project ID", () => {
it.todo("should not be able to authenticate as user");
});
describe("with project ID that doesn't exist", async () => {
backendContext.set({
projectKeys: {
projectId: "invalid",
publishableClientKey: "publish-key",
secretServerKey: "secret-key",
superSecretAdminKey: "admin-key",
}
});
it("should not have client access", async ({ expect }) => {
const response = await niceBackendFetch("/api/v1", {
accessType: "client",
});
expect(response).toMatchInlineSnapshot(`
NiceResponse {
"status": 400,
"body": {
"code": "CURRENT_PROJECT_NOT_FOUND",
"details": { "project_id": "invalid" },
"error": "The current project with ID invalid was not found. Please check the value of the x-stack-project-id header.",
},
"headers": Headers {
"x-stack-known-error": "CURRENT_PROJECT_NOT_FOUND",
<some fields may have been hidden>,
},
}
`);
});
});
describe("with a branch header that doesn't exist", async () => {
backendContext.set({
projectKeys: {
...InternalProjectKeys,
},
currentBranchId: "invalid-branch",
});
it("should not have client access", async ({ expect }) => {
const response = await niceBackendFetch("/api/v1", {
accessType: "client",
});
expect(response).toMatchInlineSnapshot(`
NiceResponse {
"status": 400,
"body": {
"code": "BRANCH_DOES_NOT_EXIST",
"details": { "branch_id": "invalid-branch" },
"error": "The branch with ID invalid-branch does not exist.",
},
"headers": Headers {
"x-stack-known-error": "BRANCH_DOES_NOT_EXIST",
<some fields may have been hidden>,
},
}
`);
});
});
describe("with project keys that don't exist", async () => {
backendContext.set({
projectKeys: {

View File

@ -5,7 +5,7 @@ import { Auth, InternalApiKey, InternalProjectKeys, Project, backendContext, nic
// TODO some of the tests here test /api/v1/projects/current, the others test /api/v1/internal/projects/current. We should split them into different test files
it("should not have have access to the project without project keys", async ({ expect }) => {
it("should not have access to the project without project keys", async ({ expect }) => {
backendContext.set({
projectKeys: 'no-project'
});