chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,348 @@
|
||||
import type { WorkerInstanceGroup, WorkloadType } from "@trigger.dev/database";
|
||||
import { WorkerInstanceGroupType } from "@trigger.dev/database";
|
||||
import { WithRunEngine } from "../baseService.server";
|
||||
import { isWorkerGroupAllowedForProject } from "./workerGroupAccess";
|
||||
import { WorkerGroupTokenService } from "./workerGroupTokenService.server";
|
||||
import { logger } from "~/services/logger.server";
|
||||
import { FEATURE_FLAG } from "~/v3/featureFlags";
|
||||
import { makeFlag, makeSetFlag } from "~/v3/featureFlags.server";
|
||||
import { isComputeRegionAccessible, resolveComputeAccess } from "~/v3/regionAccess.server";
|
||||
|
||||
export class WorkerGroupService extends WithRunEngine {
|
||||
private readonly defaultNamePrefix = "worker_group";
|
||||
|
||||
async createWorkerGroup({
|
||||
projectId,
|
||||
organizationId,
|
||||
name,
|
||||
description,
|
||||
type,
|
||||
hidden,
|
||||
workloadType,
|
||||
cloudProvider,
|
||||
location,
|
||||
staticIPs,
|
||||
enableFastPath,
|
||||
}: {
|
||||
projectId?: string;
|
||||
organizationId?: string;
|
||||
name?: string;
|
||||
description?: string;
|
||||
type?: WorkerInstanceGroupType;
|
||||
hidden?: boolean;
|
||||
workloadType?: WorkloadType;
|
||||
cloudProvider?: string;
|
||||
location?: string;
|
||||
staticIPs?: string;
|
||||
enableFastPath?: boolean;
|
||||
}) {
|
||||
if (!name) {
|
||||
name = await this.generateWorkerName({ projectId });
|
||||
}
|
||||
|
||||
const tokenService = new WorkerGroupTokenService({
|
||||
prisma: this._prisma,
|
||||
engine: this._engine,
|
||||
});
|
||||
const token = await tokenService.createToken();
|
||||
|
||||
const resolvedType =
|
||||
type ?? (projectId ? WorkerInstanceGroupType.UNMANAGED : WorkerInstanceGroupType.MANAGED);
|
||||
|
||||
const workerGroup = await this._prisma.workerInstanceGroup.create({
|
||||
data: {
|
||||
projectId,
|
||||
organizationId,
|
||||
type: resolvedType,
|
||||
masterQueue: this.generateMasterQueueName({ projectId, name }),
|
||||
tokenId: token.id,
|
||||
description,
|
||||
name,
|
||||
hidden,
|
||||
workloadType,
|
||||
cloudProvider,
|
||||
location,
|
||||
staticIPs,
|
||||
enableFastPath,
|
||||
},
|
||||
});
|
||||
|
||||
if (workerGroup.type === WorkerInstanceGroupType.MANAGED) {
|
||||
const _managedCount = await this._prisma.workerInstanceGroup.count({
|
||||
where: {
|
||||
type: WorkerInstanceGroupType.MANAGED,
|
||||
},
|
||||
});
|
||||
|
||||
const getFlag = makeFlag(this._prisma);
|
||||
const defaultWorkerInstanceGroupId = await getFlag({
|
||||
key: FEATURE_FLAG.defaultWorkerInstanceGroupId,
|
||||
});
|
||||
|
||||
// If there's no global default yet we should set it to the new worker group
|
||||
if (!defaultWorkerInstanceGroupId) {
|
||||
const setFlag = makeSetFlag(this._prisma);
|
||||
await setFlag({
|
||||
key: FEATURE_FLAG.defaultWorkerInstanceGroupId,
|
||||
value: workerGroup.id,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
return {
|
||||
workerGroup,
|
||||
token,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
This updates a single worker group.
|
||||
The name should never be updated. This would mean changing the masterQueue name which can have unexpected consequences.
|
||||
*/
|
||||
async updateWorkerGroup({
|
||||
projectId,
|
||||
workerGroupId,
|
||||
description,
|
||||
}: {
|
||||
projectId: string;
|
||||
workerGroupId: string;
|
||||
description?: string;
|
||||
}) {
|
||||
const workerGroup = await this._prisma.workerInstanceGroup.findUnique({
|
||||
where: {
|
||||
id: workerGroupId,
|
||||
projectId,
|
||||
},
|
||||
});
|
||||
|
||||
if (!workerGroup) {
|
||||
logger.error("[WorkerGroupService] No worker group found for update", {
|
||||
workerGroupId,
|
||||
description,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
await this._prisma.workerInstanceGroup.update({
|
||||
where: {
|
||||
id: workerGroup.id,
|
||||
},
|
||||
data: {
|
||||
description,
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
This lists worker groups.
|
||||
Without a project ID, only shared worker groups will be returned.
|
||||
With a project ID, in addition to all shared worker groups, ones associated with the project will also be returned.
|
||||
*/
|
||||
async listWorkerGroups({ projectId, listHidden }: { projectId?: string; listHidden?: boolean }) {
|
||||
const workerGroups = await this._prisma.workerInstanceGroup.findMany({
|
||||
where: {
|
||||
OR: [
|
||||
{
|
||||
type: WorkerInstanceGroupType.MANAGED,
|
||||
},
|
||||
{
|
||||
projectId,
|
||||
},
|
||||
],
|
||||
AND: listHidden ? [] : [{ hidden: false }],
|
||||
},
|
||||
});
|
||||
|
||||
return workerGroups;
|
||||
}
|
||||
|
||||
async deleteWorkerGroup({
|
||||
projectId,
|
||||
workerGroupId,
|
||||
}: {
|
||||
projectId: string;
|
||||
workerGroupId: string;
|
||||
}) {
|
||||
const workerGroup = await this._prisma.workerInstanceGroup.findUnique({
|
||||
where: {
|
||||
id: workerGroupId,
|
||||
},
|
||||
});
|
||||
|
||||
if (!workerGroup) {
|
||||
logger.error("[WorkerGroupService] WorkerGroup not found for deletion", {
|
||||
workerGroupId,
|
||||
projectId,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
if (workerGroup.projectId !== projectId) {
|
||||
logger.error("[WorkerGroupService] WorkerGroup does not belong to project", {
|
||||
workerGroupId,
|
||||
projectId,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
await this._prisma.workerInstanceGroup.delete({
|
||||
where: {
|
||||
id: workerGroupId,
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
async getGlobalDefaultWorkerGroup() {
|
||||
const flags = makeFlag(this._prisma);
|
||||
|
||||
const defaultWorkerInstanceGroupId = await flags({
|
||||
key: FEATURE_FLAG.defaultWorkerInstanceGroupId,
|
||||
});
|
||||
|
||||
if (!defaultWorkerInstanceGroupId) {
|
||||
logger.error("[WorkerGroupService] Default worker group not found in feature flags");
|
||||
return;
|
||||
}
|
||||
|
||||
const workerGroup = await this._prisma.workerInstanceGroup.findUnique({
|
||||
where: {
|
||||
id: defaultWorkerInstanceGroupId,
|
||||
},
|
||||
});
|
||||
|
||||
if (!workerGroup) {
|
||||
logger.error("[WorkerGroupService] Default worker group not found", {
|
||||
defaultWorkerInstanceGroupId,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
return workerGroup;
|
||||
}
|
||||
|
||||
async getDefaultWorkerGroupForProject({
|
||||
projectId,
|
||||
regionOverride,
|
||||
}: {
|
||||
projectId: string;
|
||||
regionOverride?: string;
|
||||
}): Promise<WorkerInstanceGroup | undefined> {
|
||||
const project = await this._prisma.project.findFirst({
|
||||
where: {
|
||||
id: projectId,
|
||||
},
|
||||
include: {
|
||||
defaultWorkerGroup: true,
|
||||
organization: { select: { featureFlags: true } },
|
||||
},
|
||||
});
|
||||
|
||||
if (!project) {
|
||||
throw new Error("Project not found.");
|
||||
}
|
||||
|
||||
// If they've specified a region, we need to check they have access to it
|
||||
if (regionOverride) {
|
||||
const workerGroup = await this._prisma.workerInstanceGroup.findFirst({
|
||||
where: {
|
||||
masterQueue: regionOverride,
|
||||
},
|
||||
});
|
||||
|
||||
if (!workerGroup) {
|
||||
throw new Error(`The region you specified doesn't exist ("${regionOverride}").`);
|
||||
}
|
||||
|
||||
// The masterQueue-only lookup above can resolve another project's
|
||||
// UNMANAGED group, so reject groups not usable by this project
|
||||
// (see isWorkerGroupAllowedForProject).
|
||||
if (!isWorkerGroupAllowedForProject(workerGroup, project.id)) {
|
||||
throw new Error(`The region you specified isn't available to you ("${regionOverride}").`);
|
||||
}
|
||||
|
||||
// If they're restricted, check they have access
|
||||
if (project.allowedWorkerQueues.length > 0) {
|
||||
if (project.allowedWorkerQueues.includes(workerGroup.masterQueue)) {
|
||||
return workerGroup;
|
||||
}
|
||||
|
||||
throw new Error(
|
||||
`You don't have access to this region ("${regionOverride}"). You can use the following regions: ${project.allowedWorkerQueues.join(
|
||||
", "
|
||||
)}.`
|
||||
);
|
||||
}
|
||||
|
||||
if (workerGroup.hidden) {
|
||||
throw new Error(`The region you specified isn't available to you ("${regionOverride}").`);
|
||||
}
|
||||
|
||||
if (workerGroup.workloadType === "MICROVM") {
|
||||
const hasComputeAccess = await resolveComputeAccess(
|
||||
this._prisma,
|
||||
project.organization.featureFlags
|
||||
);
|
||||
|
||||
if (!isComputeRegionAccessible(workerGroup, hasComputeAccess)) {
|
||||
throw new Error(`The region you specified isn't available to you ("${regionOverride}").`);
|
||||
}
|
||||
}
|
||||
|
||||
return workerGroup;
|
||||
}
|
||||
|
||||
if (project.defaultWorkerGroup) {
|
||||
return project.defaultWorkerGroup;
|
||||
}
|
||||
|
||||
return await this.getGlobalDefaultWorkerGroup();
|
||||
}
|
||||
|
||||
async setDefaultWorkerGroupForProject({
|
||||
projectId,
|
||||
workerGroupId,
|
||||
}: {
|
||||
projectId: string;
|
||||
workerGroupId: string;
|
||||
}) {
|
||||
const workerGroup = await this._prisma.workerInstanceGroup.findUnique({
|
||||
where: {
|
||||
id: workerGroupId,
|
||||
},
|
||||
});
|
||||
|
||||
if (!workerGroup) {
|
||||
logger.error("[WorkerGroupService] WorkerGroup not found", {
|
||||
workerGroupId,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
await this._prisma.project.update({
|
||||
where: {
|
||||
id: projectId,
|
||||
},
|
||||
data: {
|
||||
defaultWorkerGroupId: workerGroupId,
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
private async generateWorkerName({ projectId }: { projectId?: string }) {
|
||||
const workerGroups = await this._prisma.workerInstanceGroup.count({
|
||||
where: {
|
||||
projectId: projectId ?? null,
|
||||
},
|
||||
});
|
||||
|
||||
return `${this.defaultNamePrefix}_${workerGroups + 1}`;
|
||||
}
|
||||
|
||||
private generateMasterQueueName({ projectId, name }: { projectId?: string; name: string }) {
|
||||
if (!projectId) {
|
||||
return name;
|
||||
}
|
||||
|
||||
return `${projectId}-${name}`;
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user