allow admin role api keys (#2124)

* allow admin role api keys

* bump to rerun deployment

* types needs explicit export now for APIKey

* remove api_key.role, use User.role instead

* fix formatting

* formatting

* formatting

---------

Co-authored-by: Richard Kuo <rkuo@rkuo.com>

backend/danswer/db/models.py

@@ -199,6 +199,9 @@ class ApiKey(Base):
         DateTime(timezone=True), server_default=func.now()
     )
 
+    # Add this relationship to access the User object via user_id
+    user: Mapped["User"] = relationship("User", foreign_keys=[user_id])
+
 
 class Notification(Base):
     __tablename__ = "notification"

backend/ee/danswer/auth/api_key.py

@@ -5,6 +5,7 @@ from fastapi import Request
 from passlib.hash import sha256_crypt
 from pydantic import BaseModel
 
+from danswer.auth.schemas import UserRole
 from ee.danswer.configs.app_configs import API_KEY_HASH_ROUNDS
 
 
@@ -19,6 +20,7 @@ class ApiKeyDescriptor(BaseModel):
     api_key_display: str
     api_key: str | None = None  # only present on initial creation
     api_key_name: str | None = None
+    api_key_role: UserRole
 
     user_id: uuid.UUID
 

backend/ee/danswer/db/api_key.py

@@ -2,9 +2,9 @@ import uuid
 
 from fastapi_users.password import PasswordHelper
 from sqlalchemy import select
+from sqlalchemy.orm import joinedload
 from sqlalchemy.orm import Session
 
-from danswer.auth.schemas import UserRole
 from danswer.configs.constants import DANSWER_API_KEY_DUMMY_EMAIL_DOMAIN
 from danswer.configs.constants import DANSWER_API_KEY_PREFIX
 from danswer.configs.constants import UNNAMED_KEY_PLACEHOLDER
@@ -22,10 +22,15 @@ def is_api_key_email_address(email: str) -> bool:
 
 
 def fetch_api_keys(db_session: Session) -> list[ApiKeyDescriptor]:
-    api_keys = db_session.scalars(select(ApiKey)).all()
+    api_keys = (
+        db_session.scalars(select(ApiKey).options(joinedload(ApiKey.user)))
+        .unique()
+        .all()
+    )
     return [
         ApiKeyDescriptor(
             api_key_id=api_key.id,
+            api_key_role=api_key.user.role,
             api_key_display=api_key.api_key_display,
             api_key_name=api_key.name,
             user_id=api_key.user_id,
@@ -67,7 +72,7 @@ def insert_api_key(
         is_active=True,
         is_superuser=False,
         is_verified=True,
-        role=UserRole.BASIC,
+        role=api_key_args.role,
     )
     db_session.add(api_key_user_row)
 
@@ -83,6 +88,7 @@ def insert_api_key(
     db_session.commit()
     return ApiKeyDescriptor(
         api_key_id=api_key_row.id,
+        api_key_role=api_key_user_row.role,
         api_key_display=api_key_row.api_key_display,
         api_key=api_key,
         api_key_name=api_key_args.name,
@@ -106,12 +112,14 @@ def update_api_key(
 
     email_name = api_key_args.name or UNNAMED_KEY_PLACEHOLDER
     api_key_user.email = get_api_key_fake_email(email_name, str(api_key_user.id))
+    api_key_user.role = api_key_args.role
     db_session.commit()
 
     return ApiKeyDescriptor(
         api_key_id=existing_api_key.id,
         api_key_display=existing_api_key.api_key_display,
         api_key_name=api_key_args.name,
+        api_key_role=api_key_user.role,
         user_id=existing_api_key.user_id,
     )
 
@@ -122,6 +130,12 @@ def regenerate_api_key(db_session: Session, api_key_id: int) -> ApiKeyDescriptor
     if existing_api_key is None:
         raise ValueError(f"API key with id {api_key_id} does not exist")
 
+    api_key_user = db_session.scalar(
+        select(User).where(User.id == existing_api_key.user_id)  # type: ignore
+    )
+    if api_key_user is None:
+        raise RuntimeError("API Key does not have associated user.")
+
     new_api_key = generate_api_key()
     existing_api_key.hashed_api_key = hash_api_key(new_api_key)
     existing_api_key.api_key_display = build_displayable_api_key(new_api_key)
@@ -132,6 +146,7 @@ def regenerate_api_key(db_session: Session, api_key_id: int) -> ApiKeyDescriptor
         api_key_display=existing_api_key.api_key_display,
         api_key=new_api_key,
         api_key_name=existing_api_key.name,
+        api_key_role=api_key_user.role,
         user_id=existing_api_key.user_id,
     )
 

backend/ee/danswer/server/api_key/models.py

@@ -1,5 +1,8 @@
 from pydantic import BaseModel
 
+from danswer.auth.schemas import UserRole
+
 
 class APIKeyArgs(BaseModel):
     name: str | None = None
+    role: UserRole = UserRole.BASIC

web/src/app/ee/admin/api-key/DanswerApiKeyForm.tsx

@@ -1,10 +1,15 @@
 import { Form, Formik } from "formik";
 import { PopupSpec } from "@/components/admin/connectors/Popup";
-import { TextFormField } from "@/components/admin/connectors/Field";
+import {
+  BooleanFormField,
+  TextFormField,
+} from "@/components/admin/connectors/Field";
 import { createApiKey, updateApiKey } from "./lib";
 import { Modal } from "@/components/Modal";
 import { XIcon } from "@/components/icons/icons";
 import { Button, Divider, Text } from "@tremor/react";
+import { UserRole } from "@/lib/types";
+import { APIKey } from "./types";
 
 interface DanswerApiKeyFormProps {
   onClose: () => void;
@@ -42,14 +47,25 @@ export const DanswerApiKeyForm = ({
         <Formik
           initialValues={{
             name: apiKey?.api_key_name || "",
+            is_admin: apiKey?.api_key_role == "admin" ?? false,
           }}
           onSubmit={async (values, formikHelpers) => {
             formikHelpers.setSubmitting(true);
+
+            // Map the boolean to a UserRole string
+            const role: UserRole = values.is_admin ? "admin" : "basic";
+
+            // Prepare the payload with the UserRole
+            const payload = {
+              ...values,
+              role, // Assign the role directly as a UserRole type
+            };
+
             let response;
             if (isUpdate) {
-              response = await updateApiKey(apiKey.api_key_id, values);
+              response = await updateApiKey(apiKey.api_key_id, payload);
             } else {
-              response = await createApiKey(values);
+              response = await createApiKey(payload);
             }
             formikHelpers.setSubmitting(false);
             if (response.ok) {
@@ -88,6 +104,15 @@ export const DanswerApiKeyForm = ({
                 autoCompleteDisabled={true}
               />
 
+              <BooleanFormField
+                small
+                noPadding
+                alignTop
+                name="is_admin"
+                label="Is Admin?"
+                subtext="If set, this API key will have access to admin level server API's."
+              />
+
               <div className="flex">
                 <Button
                   type="submit"

web/src/app/ee/admin/api-key/page.tsx

@@ -26,6 +26,7 @@ import { Modal } from "@/components/Modal";
 import { Spinner } from "@/components/Spinner";
 import { deleteApiKey, regenerateApiKey } from "./lib";
 import { DanswerApiKeyForm } from "./DanswerApiKeyForm";
+import { APIKey } from "./types";
 
 const API_KEY_TEXT = `
 API Keys allow you to access Danswer APIs programmatically. Click the button below to generate a new API Key.
@@ -173,6 +174,7 @@ function Main() {
           <TableRow>
             <TableHeaderCell>Name</TableHeaderCell>
             <TableHeaderCell>API Key</TableHeaderCell>
+            <TableHeaderCell>Role</TableHeaderCell>
             <TableHeaderCell>Regenerate</TableHeaderCell>
             <TableHeaderCell>Delete</TableHeaderCell>
           </TableRow>
@@ -201,6 +203,9 @@ function Main() {
               <TableCell className="max-w-64">
                 {apiKey.api_key_display}
               </TableCell>
+              <TableCell className="max-w-64">
+                {apiKey.api_key_role.toUpperCase()}
+              </TableCell>
               <TableCell>
                 <div
                   className={`

web/src/app/ee/admin/api-key/types.ts

@@ -1,11 +1,15 @@
-interface APIKey {
+import { UserRole } from "@/lib/types";
+
+export interface APIKey {
   api_key_id: number;
   api_key_display: string;
   api_key: string | null;
   api_key_name: string | null;
+  api_key_role: UserRole;
   user_id: string;
 }
 
-interface APIKeyArgs {
+export interface APIKeyArgs {
   name?: string;
+  role: UserRole;
 }

web/src/lib/types.ts

@@ -14,13 +14,15 @@ export enum UserStatus {
   deactivated = "deactivated",
 }
 
+export type UserRole = "basic" | "admin";
+
 export interface User {
   id: string;
   email: string;
   is_active: string;
   is_superuser: string;
   is_verified: string;
-  role: "basic" | "admin";
+  role: UserRole;
   preferences: UserPreferences;
   status: UserStatus;
   current_token_created_at?: Date;