Add names to API Keys (#63)

This commit is contained in:
Alan Hagedorn 2024-04-13 12:08:36 -07:00 committed by Chris Weaver
parent 3466f6d3a4
commit 81e9880d9d
8 changed files with 267 additions and 41 deletions

View File

@ -18,6 +18,7 @@ class ApiKeyDescriptor(BaseModel):
api_key_id: int
api_key_display: str
api_key: str | None = None # only present on initial creation
api_key_name: str | None = None
user_id: uuid.UUID

View File

@ -12,6 +12,7 @@ from ee.danswer.auth.api_key import build_displayable_api_key
from ee.danswer.auth.api_key import generate_api_key
from ee.danswer.auth.api_key import hash_api_key
from ee.danswer.db.constants import DANSWER_API_KEY_DUMMY_EMAIL_DOMAIN
from ee.danswer.server.api_key.models import APIKeyArgs
_DANSWER_API_KEY = "danswer_api_key"
@ -28,6 +29,7 @@ def fetch_api_keys(db_session: Session) -> list[ApiKeyDescriptor]:
ApiKeyDescriptor(
api_key_id=api_key.id,
api_key_display=api_key.api_key_display,
api_key_name=api_key.name,
user_id=api_key.user_id,
)
for api_key in api_keys
@ -44,7 +46,9 @@ def fetch_user_for_api_key(hashed_api_key: str, db_session: Session) -> User | N
return db_session.scalar(select(User).where(User.id == api_key.user_id)) # type: ignore
def insert_api_key(db_session: Session, user_id: uuid.UUID | None) -> ApiKeyDescriptor:
def insert_api_key(
db_session: Session, api_key_args: APIKeyArgs, user_id: uuid.UUID | None
) -> ApiKeyDescriptor:
std_password_helper = PasswordHelper()
api_key = generate_api_key()
api_key_user_id = uuid.uuid4()
@ -62,6 +66,7 @@ def insert_api_key(db_session: Session, user_id: uuid.UUID | None) -> ApiKeyDesc
db_session.add(api_key_user_row)
api_key_row = ApiKey(
name=api_key_args.name,
hashed_api_key=hash_api_key(api_key),
api_key_display=build_displayable_api_key(api_key),
user_id=api_key_user_id,
@ -74,10 +79,29 @@ def insert_api_key(db_session: Session, user_id: uuid.UUID | None) -> ApiKeyDesc
api_key_id=api_key_row.id,
api_key_display=api_key_row.api_key_display,
api_key=api_key,
api_key_name=api_key_args.name,
user_id=api_key_user_id,
)
def update_api_key(
db_session: Session, api_key_id: int, api_key_args: APIKeyArgs
) -> ApiKeyDescriptor:
existing_api_key = db_session.scalar(select(ApiKey).where(ApiKey.id == api_key_id))
if existing_api_key is None:
raise ValueError(f"API key with id {api_key_id} does not exist")
existing_api_key.name = api_key_args.name
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,
user_id=existing_api_key.user_id,
)
def regenerate_api_key(db_session: Session, api_key_id: int) -> ApiKeyDescriptor:
"""NOTE: currently, any admin can regenerate any API key."""
existing_api_key = db_session.scalar(select(ApiKey).where(ApiKey.id == api_key_id))
@ -93,6 +117,7 @@ def regenerate_api_key(db_session: Session, api_key_id: int) -> ApiKeyDescriptor
api_key_id=existing_api_key.id,
api_key_display=existing_api_key.api_key_display,
api_key=new_api_key,
api_key_name=existing_api_key.name,
user_id=existing_api_key.user_id,
)

View File

@ -10,6 +10,8 @@ from ee.danswer.db.api_key import fetch_api_keys
from ee.danswer.db.api_key import insert_api_key
from ee.danswer.db.api_key import regenerate_api_key
from ee.danswer.db.api_key import remove_api_key
from ee.danswer.db.api_key import update_api_key
from ee.danswer.server.api_key.models import APIKeyArgs
router = APIRouter(prefix="/admin/api-key")
@ -24,13 +26,14 @@ def list_api_keys(
@router.post("")
def create_api_key(
api_key_args: APIKeyArgs,
user: db_models.User | None = Depends(current_admin_user),
db_session: Session = Depends(get_session),
) -> ApiKeyDescriptor:
return insert_api_key(db_session, user.id if user else None)
return insert_api_key(db_session, api_key_args, user.id if user else None)
@router.patch("/{api_key_id}")
@router.post("/{api_key_id}/regenerate")
def regenerate_existing_api_key(
api_key_id: int,
_: db_models.User | None = Depends(current_admin_user),
@ -39,6 +42,16 @@ def regenerate_existing_api_key(
return regenerate_api_key(db_session, api_key_id)
@router.patch("/{api_key_id}")
def update_existing_api_key(
api_key_id: int,
api_key_args: APIKeyArgs,
_: db_models.User | None = Depends(current_admin_user),
db_session: Session = Depends(get_session),
) -> ApiKeyDescriptor:
return update_api_key(db_session, api_key_id, api_key_args)
@router.delete("/{api_key_id}")
def delete_api_key(
api_key_id: int,

View File

@ -0,0 +1,5 @@
from pydantic import BaseModel
class APIKeyArgs(BaseModel):
name: str | None = None

View File

@ -0,0 +1,108 @@
import { Form, Formik } from "formik";
import { PopupSpec } from "@/components/admin/connectors/Popup";
import { 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";
interface DanswerApiKeyFormProps {
onClose: () => void;
setPopup: (popupSpec: PopupSpec | null) => void;
onCreateApiKey: (apiKey: APIKey) => void;
apiKey?: APIKey;
}
export const DanswerApiKeyForm = ({
onClose,
setPopup,
onCreateApiKey,
apiKey,
}: DanswerApiKeyFormProps) => {
const isUpdate = apiKey !== undefined;
return (
<Modal onOutsideClick={onClose} width="w-2/6">
<div className="px-8 py-6 bg-background">
<h2 className="text-xl font-bold flex">
{isUpdate ? "Update API Key" : "Create a new API Key"}
<div
onClick={onClose}
className="ml-auto hover:bg-hover p-1.5 rounded"
>
<XIcon
size={20}
className="my-auto flex flex-shrink-0 cursor-pointer"
/>
</div>
</h2>
<Divider />
<Formik
initialValues={{
name: apiKey?.api_key_name || "",
}}
onSubmit={async (values, formikHelpers) => {
formikHelpers.setSubmitting(true);
let response;
if (isUpdate) {
response = await updateApiKey(apiKey.api_key_id, values);
} else {
response = await createApiKey(values);
}
formikHelpers.setSubmitting(false);
if (response.ok) {
setPopup({
message: isUpdate
? "Successfully updated API key!"
: "Successfully created API key!",
type: "success",
});
if (!isUpdate) {
onCreateApiKey(await response.json());
}
onClose();
} else {
const responseJson = await response.json();
const errorMsg = responseJson.detail || responseJson.message;
setPopup({
message: isUpdate
? `Error updating API key - ${errorMsg}`
: `Error creating API key - ${errorMsg}`,
type: "error",
});
}
}}
>
{({ isSubmitting, values, setFieldValue }) => (
<Form>
<Text className="mb-4 text-lg">
Choose a memorable name for your API key. This is optional and
can be added or changed later!
</Text>
<TextFormField
name="name"
label="Name (optional):"
autoCompleteDisabled={true}
/>
<div className="flex">
<Button
type="submit"
size="xs"
color="green"
disabled={isSubmitting}
className="mx-auto w-64"
>
{isUpdate ? "Update!" : "Create!"}
</Button>
</div>
</Form>
)}
</Formik>
</div>
</Modal>
);
};

View File

@ -0,0 +1,37 @@
export const createApiKey = async (apiKeyArgs: APIKeyArgs) => {
return fetch("/api/admin/api-key", {
method: "POST",
headers: {
"Content-Type": "application/json",
},
body: JSON.stringify(apiKeyArgs),
});
};
export const regenerateApiKey = async (apiKey: APIKey) => {
return fetch(`/api/admin/api-key/${apiKey.api_key_id}/regenerate`, {
method: "POST",
headers: {
"Content-Type": "application/json",
},
});
};
export const updateApiKey = async (
apiKeyId: number,
apiKeyArgs: APIKeyArgs
) => {
return fetch(`/api/admin/api-key/${apiKeyId}`, {
method: "PATCH",
headers: {
"Content-Type": "application/json",
},
body: JSON.stringify(apiKeyArgs),
});
};
export const deleteApiKey = async (apiKeyId: number) => {
return fetch(`/api/admin/api-key/${apiKeyId}`, {
method: "DELETE",
});
};

View File

@ -21,16 +21,11 @@ import { usePopup } from "@/components/admin/connectors/Popup";
import { useState } from "react";
import { Table } from "@tremor/react";
import { DeleteButton } from "@/components/DeleteButton";
import { FiCopy, FiRefreshCw, FiX } from "react-icons/fi";
import { FiCopy, FiEdit, FiRefreshCw, FiX } from "react-icons/fi";
import { Modal } from "@/components/Modal";
import { Spinner } from "@/components/Spinner";
interface APIKey {
api_key_id: number;
api_key_display: string;
api_key: string | null; // only present on initial creation
user_id: string;
}
import { deleteApiKey, regenerateApiKey } from "./lib";
import { DanswerApiKeyForm } from "./DanswerApiKeyForm";
const API_KEY_TEXT = `
API Keys allow you to access Danswer APIs programmatically. Click the button below to generate a new API Key.
@ -97,6 +92,13 @@ function Main() {
const [fullApiKey, setFullApiKey] = useState<string | null>(null);
const [keyIsGenerating, setKeyIsGenerating] = useState(false);
const [showCreateUpdateForm, setShowCreateUpdateForm] = useState(false);
const [selectedApiKey, setSelectedApiKey] = useState<APIKey | undefined>();
const handleEdit = (apiKey: APIKey) => {
setSelectedApiKey(apiKey);
setShowCreateUpdateForm(true);
};
if (isLoading) {
return <ThreeDotsLoader />;
@ -116,24 +118,7 @@ function Main() {
color="green"
size="xs"
className="mt-3"
onClick={async () => {
setKeyIsGenerating(true);
const response = await fetch("/api/admin/api-key", {
method: "POST",
});
setKeyIsGenerating(false);
if (!response.ok) {
const errorMsg = await response.text();
setPopup({
type: "error",
message: `Failed to create API Key: ${errorMsg}`,
});
return;
}
const newKey = (await response.json()) as APIKey;
setFullApiKey(newKey.api_key);
mutate("/api/admin/api-key");
}}
onClick={() => setShowCreateUpdateForm(true)}
>
Create API Key
</Button>
@ -145,12 +130,29 @@ function Main() {
{popup}
<Text>{API_KEY_TEXT}</Text>
{newApiKeyButton}
{showCreateUpdateForm && (
<DanswerApiKeyForm
onCreateApiKey={(apiKey) => {
setFullApiKey(apiKey.api_key);
}}
onClose={() => {
setShowCreateUpdateForm(false);
setSelectedApiKey(undefined);
mutate("/api/admin/api-key");
}}
setPopup={setPopup}
apiKey={selectedApiKey}
/>
)}
</div>
);
}
return (
<div>
{popup}
{fullApiKey && (
<NewApiKeyModal
apiKey={fullApiKey}
@ -169,6 +171,7 @@ function Main() {
<Table className="overflow-visible">
<TableHead>
<TableRow>
<TableHeaderCell>Name</TableHeaderCell>
<TableHeaderCell>API Key</TableHeaderCell>
<TableHeaderCell>Regenerate</TableHeaderCell>
<TableHeaderCell>Delete</TableHeaderCell>
@ -177,6 +180,24 @@ function Main() {
<TableBody>
{apiKeys.map((apiKey) => (
<TableRow key={apiKey.api_key_id}>
<TableCell>
<div
className={`
my-auto
flex
mb-1
w-fit
hover:bg-hover cursor-pointer
p-2
rounded-lg
border-border
text-sm`}
onClick={() => handleEdit(apiKey)}
>
<FiEdit className="my-auto mr-2" />
{apiKey.api_key_name || <i>null</i>}
</div>
</TableCell>
<TableCell className="max-w-64">
{apiKey.api_key_display}
</TableCell>
@ -194,12 +215,7 @@ function Main() {
text-sm`}
onClick={async () => {
setKeyIsGenerating(true);
const response = await fetch(
`/api/admin/api-key/${apiKey.api_key_id}`,
{
method: "PATCH",
}
);
const response = await regenerateApiKey(apiKey);
setKeyIsGenerating(false);
if (!response.ok) {
const errorMsg = await response.text();
@ -221,12 +237,7 @@ function Main() {
<TableCell>
<DeleteButton
onClick={async () => {
const response = await fetch(
`/api/admin/api-key/${apiKey.api_key_id}`,
{
method: "DELETE",
}
);
const response = await deleteApiKey(apiKey.api_key_id);
if (!response.ok) {
const errorMsg = await response.text();
setPopup({
@ -243,6 +254,21 @@ function Main() {
))}
</TableBody>
</Table>
{showCreateUpdateForm && (
<DanswerApiKeyForm
onCreateApiKey={(apiKey) => {
setFullApiKey(apiKey.api_key);
}}
onClose={() => {
setShowCreateUpdateForm(false);
setSelectedApiKey(undefined);
mutate("/api/admin/api-key");
}}
setPopup={setPopup}
apiKey={selectedApiKey}
/>
)}
</div>
);
}

View File

@ -0,0 +1,11 @@
interface APIKey {
api_key_id: number;
api_key_display: string;
api_key: string | null;
api_key_name: string | null;
user_id: string;
}
interface APIKeyArgs {
name?: string;
}