unkeyed · ogzhanolguncu · Dec 20, 2024 · chronark · Dec 20, 2024 · ogzhanolguncu
@@ -12,7 +12,7 @@ import { FormField } from "@/components/ui/form";
 import { Textarea } from "@/components/ui/textarea";
 import { toast } from "@/components/ui/toaster";
 import { trpc } from "@/lib/trpc/client";
-import { cn } from "@/lib/utils";
+import { cn, getFlag } from "@/lib/utils";
 import { zodResolver } from "@hookform/resolvers/zod";
 import type { Workspace } from "@unkey/db";
 import { Button } from "@unkey/ui";
@@ -41,7 +41,7 @@ type Props = {
 
 export const UpdateIpWhitelist: React.FC<Props> = ({ api, workspace }) => {
   const router = useRouter();
-  const isEnabled = workspace.features.ipWhitelist;
+  const isEnabled = getFlag(workspace, "ipWhitelist");
 
   const form = useForm<z.infer<typeof formSchema>>({
     resolver: zodResolver(formSchema),

@@ -7,6 +7,7 @@ import { PageContent } from "@/components/page-content";
 import { Table, TableBody, TableHead, TableHeader, TableRow } from "@/components/ui/table";
 import { getTenantId } from "@/lib/auth";
 import { db } from "@/lib/db";
+import { getFlag } from "@/lib/utils";
 import { Fingerprint } from "@unkey/icons";
 import { Loader2 } from "lucide-react";
 import { unstable_cache as cache } from "next/cache";
@@ -34,7 +35,7 @@ export default async function Page(props: Props) {
     return redirect("/auth/sign-in");
   }
 
-  if (!workspace.betaFeatures.identities) {
+  if (!getFlag(workspace, "identities")) {
     return <OptIn title="Identities" description="Identities are in beta" feature="identities" />;
   }
 

@@ -3,6 +3,7 @@
 import { getTenantId } from "@/lib/auth";
 import { clickhouse } from "@/lib/clickhouse";
 import { db } from "@/lib/db";
+import { getFlag } from "@/lib/utils";
 import { notFound } from "next/navigation";
 import { createSearchParamsCache } from "nuqs/server";
 import { DEFAULT_LOGS_FETCH_COUNT } from "./constants";
@@ -25,7 +26,7 @@ export default async function Page({
       and(eq(table.tenantId, tenantId), isNull(table.deletedAt)),
   });
 
-  if (!workspace?.betaFeatures.logsPage) {
+  if (!workspace || !getFlag(workspace, "logsPage")) {
     return notFound();
   }
 

@@ -7,6 +7,7 @@ import { PageContent } from "@/components/page-content";
 import { Badge } from "@/components/ui/badge";
 import { getTenantId } from "@/lib/auth";
 import { db } from "@/lib/db";
+import { getFlag } from "@/lib/utils";
 import { Gauge } from "@unkey/icons";
 import { Scan } from "lucide-react";
 import { notFound } from "next/navigation";
@@ -89,7 +90,11 @@ export default async function OverridePage(props: Props) {
             actions={[
               <Badge variant="secondary" className="h-8">
                 {Intl.NumberFormat().format(namespace.overrides.length)} /{" "}
-                {Intl.NumberFormat().format(namespace.workspace.features.ratelimitOverrides ?? 5)}{" "}
+                {Intl.NumberFormat().format(
+                  getFlag(namespace.workspace, "ratelimitOverrides", {
+                    devModeDefault: 5,
+                  }) ?? 5,
+                )}{" "}
                 used{" "}
               </Badge>,
             ]}

@@ -13,7 +13,7 @@ import {
   ShieldCheck,
   TableProperties,
 } from "lucide-react";
-import { cn } from "../../lib/utils";
+import { cn, getFlag } from "../../lib/utils";
 
 type NavItem = {
   disabled?: boolean;
@@ -91,23 +91,23 @@ export const createWorkspaceNavigation = (
       href: "/monitors/verifications",
       label: "Monitors",
       active: segments.at(0) === "verifications",
-      hidden: !workspace.features.webhooks,
+      hidden: !getFlag(workspace, "webhooks"),
     },
     {
       icon: TableProperties,
       href: "/logs",
       label: "Logs",
       active: segments.at(0) === "logs",
       tag: <Tag label="Beta" />,
-      hidden: !workspace.betaFeatures.logsPage,
+      hidden: !getFlag(workspace, "logsPage"),
     },
     {
       icon: Crown,
       href: "/success",
       label: "Success",
       active: segments.at(0) === "success",
       tag: <Tag label="Internal" />,
-      hidden: !workspace.features.successPage,
+      hidden: !getFlag(workspace, "successPage"),
     },
     {
       icon: DatabaseZap,
@@ -121,7 +121,7 @@ export const createWorkspaceNavigation = (
       href: "/identities",
       label: "Identities",
       active: segments.at(0) === "identities",
-      hidden: !workspace.betaFeatures.identities,
+      hidden: !getFlag(workspace, "identities"),
     },
     {
       icon: Settings2,

@@ -4,6 +4,7 @@ import { z } from "zod";
 import { insertAuditLogs } from "@/lib/audit";
 import { db, eq, schema } from "@/lib/db";
 
+import { getFlag } from "@/lib/utils";
 import { auth, t } from "../../trpc";
 
 export const updateApiIpWhitelist = t.procedure
@@ -58,7 +59,8 @@ export const updateApiIpWhitelist = t.procedure
       });
     }
 
-    if (!api.workspace.features.ipWhitelist) {
+    const hasIpWhitelistFeatureFlag = getFlag(api.workspace, "ipWhitelist");
+    if (!hasIpWhitelistFeatureFlag) {
       throw new TRPCError({
         code: "FORBIDDEN",
         message:

@@ -1,6 +1,7 @@
 import { DEFAULT_FETCH_COUNT } from "@/app/(app)/audit/components/table/constants";
 import { type Workspace, db } from "@/lib/db";
 import { rateLimitedProcedure, ratelimit } from "@/lib/trpc/ratelimitProcedure";
+import { getFlag } from "@/lib/utils";
 import { type User, clerkClient } from "@clerk/nextjs/server";
 import { TRPCError } from "@trpc/server";
 import type { SelectAuditLog, SelectAuditLogTarget } from "@unkey/db/src/schema";
@@ -124,8 +125,11 @@ export const queryAuditLogs = async (options: QueryOptions, workspace: Workspace
     limit = DEFAULT_FETCH_COUNT,
   } = options;
 
-  const retentionDays =
-    workspace.features.auditLogRetentionDays ?? workspace.plan === "free" ? 30 : 90;
+  const auditLogRetentionDays = getFlag(workspace, "auditLogRetentionDays", {
+    devModeDefault: 90,
+  });
+  const retentionDays = auditLogRetentionDays ?? workspace.plan === "free" ? 30 : 90;
+
   const retentionCutoffUnixMilli = Date.now() - retentionDays * 24 * 60 * 60 * 1000;
 
   return db.query.auditLogBucket.findFirst({

@@ -3,6 +3,7 @@ import { z } from "zod";
 
 import { insertAuditLogs } from "@/lib/audit";
 import { and, db, eq, isNull, schema, sql } from "@/lib/db";
+import { getFlag } from "@/lib/utils";
 import { newId } from "@unkey/id";
 import { auth, t } from "../../trpc";
 export const createOverride = t.procedure
@@ -60,9 +61,10 @@ export const createOverride = t.procedure
           )
           .then((res) => Number(res.at(0)?.count ?? 0));
         const max =
-          typeof namespace.workspace.features.ratelimitOverrides === "number"
-            ? namespace.workspace.features.ratelimitOverrides
-            : 5;
+          getFlag(namespace.workspace, "ratelimitOverrides", {
+            devModeDefault: 5,
+          }) ?? 5;
+
         if (existing >= max) {
           throw new TRPCError({
             code: "FORBIDDEN",

@@ -1,3 +1,4 @@
+import type { Workspace } from "@unkey/db/src/types";
 import { type ClassValue, clsx } from "clsx";
 import { twMerge } from "tailwind-merge";
 
@@ -128,3 +129,64 @@ export function throttle<T extends (...args: any[]) => any>(
 
   return throttled;
 }
+
+type WorkspaceFeatures = Pick<Workspace, "features" | "betaFeatures">;
+
+type ConfigObject = WorkspaceFeatures["betaFeatures"] & WorkspaceFeatures["features"];
+
+type FlagValue<T extends keyof ConfigObject> = NonNullable<ConfigObject[T]>;
+
+/**
+ * Checks if a workspace has access to a specific feature or beta feature.
+ * In development environment, returns devModeDefault value or true if not specified.
+ * In production, returns the feature value if set, otherwise returns undefined.
+ *
+ * @param workspace - The workspace to check access for
+ * @param flagName - The name of the feature to check
+ * @param options - Configuration options
+ * @param options.devModeDefault - Value to return in development environment. Defaults to true if not specified.
+ *                                For boolean flags, you don't need to specify true as it's the default.
+ * @returns The feature value (boolean | number | string) or undefined if not set in production
+ *
+ * @example
+ * ```typescript
+ * // Check if workspace has access to logs page (boolean flag)
+ * const hasLogsAccess = getFlag(workspace, "logsPage");
+ * if (!hasLogsAccess) {
+ *   return notFound();
+ * }
+ *
+ * // Check feature with numeric value
+ * const userLimit = getFlag(workspace, "userLimit", {
+ *   devModeDefault: 1000
+ * });
+ *
+ * // Check feature with string value
+ * const tier = getFlag(workspace, "serviceTier", {
+ *   devModeDefault: "premium"
+ * });
+ * ```
+ */
+export function getFlag<TFlagName extends keyof ConfigObject>(
+  workspace: Partial<WorkspaceFeatures>,
+  flagName: TFlagName,
+  options: { devModeDefault?: FlagValue<TFlagName> } = {},
+): FlagValue<TFlagName> | undefined {
+  if (process.env.NODE_ENV === "development") {
+    return options.devModeDefault ?? (true as FlagValue<TFlagName>);
+  }
+
+  if (!workspace) {
+    throw new Error(
+      "Cannot get feature flag: No workspace found in database. Please verify workspace exists in the database or create a new workspace record.",
+    );
+  }
+
+  const betaFeature = workspace.betaFeatures?.[flagName as keyof Workspace["betaFeatures"]];
+  if (betaFeature != null) {
+    return betaFeature as FlagValue<TFlagName> | undefined;
+  }
+
+  const feature = workspace.features?.[flagName as keyof Workspace["features"]];
+  return feature as FlagValue<TFlagName> | undefined;
+}