fix: transaction deadlock in trpc

.github/workflows/job_test_api_canary.yaml

@@ -24,7 +24,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        shard: ["1/9", "2/9", "3/9", "4/9", "5/9", "6/9","7/9", "8/9","9/9"]
+        shard: ["1/9", "2/9", "3/9", "4/9", "5/9", "6/9", "7/9", "8/9", "9/9"]
 
     steps:
       - uses: actions/checkout@v4

.github/workflows/job_test_api_local.yaml

@@ -13,12 +13,9 @@ jobs:
       - name: Delete huge unnecessary tools folder
         run: rm -rf /opt/hostedtoolcache
 
-
       - name: Run containers
         run: docker compose -f ./deployment/docker-compose.yaml up -d
 
-
-
       - name: Install
         uses: ./.github/actions/install
         with:
@@ -28,7 +25,6 @@ jobs:
       - name: Build
         run: pnpm turbo run build --filter=./apps/api...
 
-
       - name: Load Schema into MySQL
         run: pnpm drizzle-kit push
         working-directory: internal/db
@@ -50,7 +46,7 @@ jobs:
           DATABASE_HOST: localhost:3900
           DATABASE_USERNAME: unkey
           DATABASE_PASSWORD: password
-          CLICKHOUSE_URL:  http://default:password@localhost:8123
+          CLICKHOUSE_URL: http://default:password@localhost:8123
           TEST_LOCAL: true
 
       - name: Dump logs

.github/workflows/job_test_api_staging.yaml

@@ -24,7 +24,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        shard: ["1/9", "2/9", "3/9", "4/9", "5/9", "6/9","7/9", "8/9", "9/9"]
+        shard: ["1/9", "2/9", "3/9", "4/9", "5/9", "6/9", "7/9", "8/9", "9/9"]
 
     steps:
       - uses: actions/checkout@v4
@@ -34,8 +34,6 @@ jobs:
       - name: Wake ClickHouse
         run: curl -X GET ${{ secrets.CLICKHOUSE_URL }}/ping
 
-
-
       - name: Install
         uses: ./.github/actions/install
         with:

.github/workflows/pr.yaml

@@ -46,5 +46,5 @@ jobs:
   #       working-directory: apps/agent
 
   # test_go_api_local:
-  #   name: Test GO API Local
+  #   name: Test Go API Local
   #   uses: ./.github/workflows/job_test_go_api_local.yaml

apps/dashboard/app/(app)/logs/components/controls/components/logs-filters/components/paths-filter.tsx

@@ -3,25 +3,21 @@ import { useFilters } from "@/app/(app)/logs/hooks/use-filters";
 import { Checkbox } from "@/components/ui/checkbox";
 import { trpc } from "@/lib/trpc/client";
 import { Button } from "@unkey/ui";
-import { useCallback, useMemo } from "react";
+import { useCallback } from "react";
 import { useCheckboxState } from "./hooks/use-checkbox-state";
 
 export const PathsFilter = () => {
-  const dateNow = useMemo(() => Date.now(), []);
-  const { data: paths, isLoading } = trpc.logs.queryDistinctPaths.useQuery(
-    { currentDate: dateNow },
-    {
-      select(paths) {
-        return paths
-          ? paths.map((path, index) => ({
-              id: index + 1,
-              path,
-              checked: false,
-            }))
-          : [];
-      },
+  const { data: paths, isLoading } = trpc.logs.queryDistinctPaths.useQuery(undefined, {
+    select(paths) {
+      return paths
+        ? paths.map((path, index) => ({
+            id: index + 1,
+            path,
+            checked: false,
+          }))
+        : [];
     },
-  );
+  });
   const { filters, updateFilters } = useFilters();
 
   const { checkboxes, handleCheckboxChange, handleSelectAll, handleKeyDown } = useCheckboxState({

apps/dashboard/app/(app)/logs/components/controls/components/logs-search/index.tsx

@@ -50,10 +50,7 @@ export const LogsSearch = () => {
     const query = search.trim();
     if (query) {
       try {
-        await queryLLMForStructuredOutput.mutateAsync({
-          query,
-          timestamp: Date.now(),
-        });
+        await queryLLMForStructuredOutput.mutateAsync(query);
       } catch (error) {
         console.error("Search failed:", error);
       }

apps/dashboard/app/(app)/settings/billing/stripe/success/page.tsx

@@ -29,6 +29,11 @@ export default async function StripeSuccess(props: Props) {
   const ws = await db.query.workspaces.findFirst({
     where: (table, { and, eq, isNull }) =>
       and(eq(table.tenantId, tenantId), isNull(table.deletedAt)),
+    with: {
+      auditLogBuckets: {
+        where: (table, { eq }) => eq(table.name, "unkey_mutations"),
+      },
+    },
   });
   if (!ws) {
     return redirect("/new");
@@ -99,7 +104,7 @@ export default async function StripeSuccess(props: Props) {
       .where(eq(schema.workspaces.id, ws.id));
 
     if (isUpgradingPlan) {
-      await insertAuditLogs(tx, {
+      await insertAuditLogs(tx, ws.auditLogBuckets[0].id, {
         workspaceId: ws.id,
         actor: { type: "user", id: user.id },
         event: "workspace.update",

apps/dashboard/app/new/create-ratelimit.tsx

@@ -4,11 +4,15 @@ import { getTenantId } from "@/lib/auth";
 import { router } from "@/lib/trpc/routers";
 import { auth } from "@clerk/nextjs";
 import { createCallerFactory } from "@trpc/server";
+import type { AuditLogBucket, Workspace } from "@unkey/db";
 import { Button } from "@unkey/ui";
 import { GlobeLock } from "lucide-react";
 import Link from "next/link";
 
-export const CreateRatelimit: React.FC = async () => {
+type Props = {
+  workspace: Workspace & { auditLogBucket: AuditLogBucket };
+};
+export const CreateRatelimit: React.FC<Props> = async (props) => {
   const { sessionClaims, userId } = auth();
   if (!userId) {
     return null;
@@ -20,6 +24,7 @@ export const CreateRatelimit: React.FC = async () => {
     user: {
       id: userId,
     },
+    workspace: props.workspace,
     tenant: {
       id: tenantId,
       role: "",
@@ -30,10 +35,6 @@ export const CreateRatelimit: React.FC = async () => {
     },
   });
 
-  await trpc.workspace.optIntoBeta({
-    feature: "ratelimit",
-  });
-
   const rootKey = await trpc.rootKey.create({
     name: "onboarding",
     permissions: ["ratelimit.*.create_namespace", "ratelimit.*.limit"],

apps/dashboard/app/new/page.tsx

@@ -162,6 +162,11 @@ export default async function (props: Props) {
     const workspace = await db.query.workspaces.findFirst({
       where: (table, { and, eq, isNull }) =>
         and(eq(table.id, props.searchParams.workspaceId!), isNull(table.deletedAt)),
+      with: {
+        auditLogBuckets: {
+          where: (table, { eq }) => eq(table.name, "unkey_mutations"),
+        },
+      },
     });
     if (!workspace) {
       return redirect("/new");
@@ -184,7 +189,9 @@ export default async function (props: Props) {
 
         <Separator className="my-8" />
 
-        <CreateRatelimit />
+        <CreateRatelimit
+          workspace={{ ...workspace, auditLogBucket: workspace.auditLogBuckets[0] }}
+        />
       </div>
     );
   }
@@ -210,7 +217,17 @@ export default async function (props: Props) {
           subscriptions: null,
           createdAt: new Date(),
         });
-        await insertAuditLogs(tx, {
+
+        const bucketId = newId("auditLogBucket");
+        await tx.insert(schema.auditLogBucket).values({
+          id: bucketId,
+          workspaceId,
+          name: "unkey_mutations",
+          retentionDays: 30,
+          deleteProtection: true,
+        });
+
+        await insertAuditLogs(tx, bucketId, {
           workspaceId: workspaceId,
           event: "workspace.create",
           actor: {

apps/dashboard/lib/audit.ts

@@ -91,14 +91,9 @@ export type UnkeyAuditLog = {
   };
 };
 
-const BUCKET_NAME = "unkey_mutations";
-
-type Key = `${string}::${string}`;
-type BucketId = string;
-const bucketCache = new Map<Key, BucketId>();
-
 export async function insertAuditLogs(
   db: Transaction | Database,
+  bucketId: string,
   logOrLogs: MaybeArray<UnkeyAuditLog>,
 ) {
   const logs = Array.isArray(logOrLogs) ? logOrLogs : [logOrLogs];
@@ -108,35 +103,6 @@ export async function insertAuditLogs(
   }
 
   for (const log of logs) {
-    // 1. Get the bucketId or create one if necessary
-    const key: Key = `${log.workspaceId}::${BUCKET_NAME}`;
-    let bucketId = "";
-    const cachedBucketId = bucketCache.get(key);
-    if (cachedBucketId) {
-      bucketId = cachedBucketId;
-    } else {
-      const bucket = await db.query.auditLogBucket.findFirst({
-        where: (table, { eq, and }) =>
-          and(eq(table.workspaceId, log.workspaceId), eq(table.name, BUCKET_NAME)),
-        columns: {
-          id: true,
-        },
-      });
-      if (bucket) {
-        bucketId = bucket.id;
-      } else {
-        bucketId = newId("auditLogBucket");
-        await db.insert(schema.auditLogBucket).values({
-          id: bucketId,
-          workspaceId: log.workspaceId,
-          name: BUCKET_NAME,
-        });
-      }
-    }
-    bucketCache.set(key, bucketId);
-
-    // 2. Insert the log
-
     const auditLogId = newId("auditLog");
     await db.insert(schema.auditLog).values({
       id: auditLogId,

apps/dashboard/lib/trpc/context.ts

@@ -2,17 +2,57 @@ import type { inferAsyncReturnType } from "@trpc/server";
 import type { FetchCreateContextFnOptions } from "@trpc/server/adapters/fetch";
 
 import { getAuth } from "@clerk/nextjs/server";
+import { newId } from "@unkey/id";
+import { type AuditLogBucket, type Workspace, db, schema } from "../db";
 
 export async function createContext({ req }: FetchCreateContextFnOptions) {
   const { userId, orgId, orgRole } = getAuth(req as any);
 
+  let ws: (Workspace & { auditLogBucket: AuditLogBucket }) | undefined;
+  const tenantId = orgId ?? userId;
+  if (tenantId) {
+    await db.transaction(async (tx) => {
+      const res = await tx.query.workspaces.findFirst({
+        where: (table, { eq, and, isNull }) =>
+          and(eq(table.tenantId, tenantId), isNull(table.deletedAt)),
+        with: {
+          auditLogBuckets: {
+            where: (table, { eq }) => eq(table.name, "unkey_mutations"),
+          },
+        },
+      });
+      if (res) {
+        let auditLogBucket = res.auditLogBuckets.at(0);
+        // @ts-expect-error it should be undefined
+        delete res.auditLogBuckets; // we don't need to pollute or context
+        if (!auditLogBucket) {
+          auditLogBucket = {
+            id: newId("auditLogBucket"),
+            name: "unkey_mutations",
+            createdAt: Date.now(),
+            deleteProtection: true,
+            workspaceId: res.id,
+            retentionDays: 30,
+            updatedAt: null,
+          };
+          await tx.insert(schema.auditLogBucket).values(auditLogBucket);
+        }
+        ws = {
+          ...res,
+          auditLogBucket,
+        };
+      }
+    });
+  }
+
   return {
     req,
     audit: {
       userAgent: req.headers.get("user-agent") ?? undefined,
       location: req.headers.get("x-forwarded-for") ?? process.env.VERCEL_REGION ?? "unknown",
     },
     user: userId ? { id: userId } : null,
+    workspace: ws,
     tenant:
       orgId && orgRole
         ? {

apps/dashboard/lib/trpc/routers/api/create.ts

@@ -12,34 +12,16 @@ export const createApi = t.procedure
     z.object({
       name: z
         .string()
-        .min(3, "workspace names must contain at least 3 characters")
-        .max(50, "workspace names must contain at most 50 characters"),
+        .min(3, "API names must contain at least 3 characters")
+        .max(50, "API names must contain at most 50 characters"),
     }),
   )
   .mutation(async ({ input, ctx }) => {
-    const ws = await db.query.workspaces
-      .findFirst({
-        where: (table, { and, eq, isNull }) =>
-          and(eq(table.tenantId, ctx.tenant.id), isNull(table.deletedAt)),
-      })
-      .catch((_err) => {
-        throw new TRPCError({
-          code: "INTERNAL_SERVER_ERROR",
-          message: "We are unable to create an API. Please try again or contact support@unkey.dev",
-        });
-      });
-    if (!ws) {
-      throw new TRPCError({
-        code: "NOT_FOUND",
-        message: "The workspace does not exist.",
-      });
-    }
-
     const keyAuthId = newId("keyAuth");
     try {
       await db.insert(schema.keyAuth).values({
         id: keyAuthId,
-        workspaceId: ws.id,
+        workspaceId: ctx.workspace.id,
         createdAt: new Date(),
       });
     } catch (_err) {
@@ -58,7 +40,7 @@ export const createApi = t.procedure
           .values({
             id: apiId,
             name: input.name,
-            workspaceId: ws.id,
+            workspaceId: ctx.workspace.id,
             keyAuthId,
             authType: "key",
             ipWhitelist: null,
@@ -72,8 +54,8 @@ export const createApi = t.procedure
             });
           });
 
-        await insertAuditLogs(tx, {
-          workspaceId: ws.id,
+        await insertAuditLogs(tx, ctx.workspace.auditLogBucket.id, {
+          workspaceId: ctx.workspace.id,
           actor: {
             type: "user",
             id: ctx.user.id,