feat(api): long lived api keys (#658)

* adds auth endpoints for creating, listing, updating, and revoking API keys
* adds api keys table and policies for reading existing tables with api keys
* adds rpc for insert api keys into the api keys table
* adds python crud operations for api keys
* adds integration tests for api keys

packages/api/supabase/migrations/20240618163044_v0.9.0_api_keys.sql

@@ -0,0 +1,175 @@
+create extension if not exists pgcrypto;
+
+-- Initialize api_keys table
+create table api_keys (
+    name text,
+    id uuid primary key default uuid_generate_v4(),
+    user_id uuid references auth.users not null,
+    api_key_hash text not null unique,
+    created_at bigint default extract(epoch from now()) not null,
+    expires_at bigint default null,
+    checksum text not null
+);
+
+alter table api_keys enable row level security;
+
+-- Hash the api key and store it in the table
+create or replace function insert_api_key(
+    p_name text,
+    p_user_id uuid,
+    p_api_key text,
+    p_checksum text,
+    p_expires_at bigint default null
+) returns table (
+    name text,
+    id uuid,
+    created_at bigint,
+    expires_at bigint,
+    checksum text
+) language plpgsql as $$
+declare
+    v_name text;
+    v_id uuid;
+    v_created_at bigint;
+    v_expires_at bigint;
+    v_checksum text;
+    v_hash text;
+begin
+    -- Calculate the one-way hash of the api key
+    v_hash := extensions.crypt(p_api_key, extensions.gen_salt('bf'));
+
+    insert into api_keys (name, user_id, api_key_hash, expires_at, checksum)
+    values (p_name, p_user_id, v_hash, p_expires_at, p_checksum)
+    returning api_keys.name, api_keys.id, api_keys.created_at, api_keys.expires_at, api_keys.checksum
+    into v_name, v_id, v_created_at, v_expires_at, v_checksum;
+
+    return query select v_name, v_id, v_created_at, v_expires_at, v_checksum;
+end;
+$$;
+
+create policy "Read only if API key matches and is current" ON api_keys for
+    select using (
+        api_key_hash = crypt(current_setting('request.headers')::json->>'x-custom-api-key', api_key_hash)
+        and (expires_at is null or expires_at > extract(epoch from now()))
+    );
+
+create policy "Individuals can crud their own api_keys." on api_keys for
+    all using (auth.uid() = user_id);
+
+-- API Key Policies
+
+create policy "Individuals can CRUD their own assistant_objects via API key."
+    on assistant_objects for all
+    to anon
+    using
+    (
+        exists (
+            select 1
+            from api_keys
+            where api_keys.api_key_hash = crypt(current_setting('request.headers')::json->>'x-custom-api-key', api_keys.api_key_hash)
+            and api_keys.user_id = assistant_objects.user_id
+        )
+    );
+
+create policy "Individuals can CRUD their own thread_objects via API key."
+    on thread_objects for all
+    to anon
+    using
+    (
+        exists (
+            select 1
+            from api_keys
+            where api_keys.api_key_hash = crypt(current_setting('request.headers')::json->>'x-custom-api-key', api_keys.api_key_hash)
+            and api_keys.user_id = thread_objects.user_id
+        )
+    );
+
+create policy "Individuals can CRUD their own message_objects via API key."
+    on message_objects for all
+    to anon
+    using
+    (
+        exists (
+            select 1
+            from api_keys
+            where api_keys.api_key_hash = crypt(current_setting('request.headers')::json->>'x-custom-api-key', api_keys.api_key_hash)
+            and api_keys.user_id = message_objects.user_id
+        )
+    );
+
+create policy "Individuals can CRUD their own file_objects via API key."
+    on file_objects for all
+    to anon
+    using
+    (
+        exists (
+            select 1
+            from api_keys
+            where api_keys.api_key_hash = crypt(current_setting('request.headers')::json->>'x-custom-api-key', api_keys.api_key_hash)
+            and api_keys.user_id = file_objects.user_id
+        )
+    );
+
+create policy "Individuals can CRUD file_bucket via API key."
+    on storage.buckets for all
+    to anon
+    using
+    (
+        exists (
+            select 1
+            from api_keys
+            where api_keys.api_key_hash = crypt(current_setting('request.headers')::json->>'x-custom-api-key', api_keys.api_key_hash)
+        )
+    );
+
+create policy "Individuals can CRUD their own run_objects via API key."
+    on run_objects for all
+    to anon
+    using
+    (
+        exists (
+            select 1
+            from api_keys
+            where api_keys.api_key_hash = crypt(current_setting('request.headers')::json->>'x-custom-api-key', api_keys.api_key_hash)
+            and api_keys.user_id = run_objects.user_id
+        )
+    );
+
+create policy "Individuals can CRUD their own vector_store via API key."
+    on vector_store for all
+    to anon
+    using
+    (
+        exists (
+            select 1
+            from api_keys
+            where api_keys.api_key_hash = crypt(current_setting('request.headers')::json->>'x-custom-api-key', api_keys.api_key_hash)
+            and api_keys.user_id = vector_store.user_id
+        )
+    );
+
+create policy "Individuals can CRUD their own vector_store_file via API key."
+    on vector_store_file for all
+    to anon
+    using
+    (
+        exists (
+            select 1
+            from api_keys
+            where api_keys.api_key_hash = crypt(current_setting('request.headers')::json->>'x-custom-api-key', api_keys.api_key_hash)
+            and api_keys.user_id = vector_store_file.user_id
+        )
+    );
+
+create policy "Individuals can CRUD their own vector_content via API key."
+    on vector_content for all
+    to anon
+    using
+    (
+        exists (
+            select 1
+            from api_keys
+            where api_keys.api_key_hash = crypt(current_setting('request.headers')::json->>'x-custom-api-key', api_keys.api_key_hash)
+            and api_keys.user_id = vector_content.user_id
+        )
+    );

src/leapfrogai_api/backend/rag/index.py

@@ -11,7 +11,7 @@
 from openai.types.beta.vector_store import FileCounts, VectorStore
 from openai.types.beta.vector_stores import VectorStoreFile
 from openai.types.beta.vector_stores.vector_store_file import LastError
-from supabase_py_async import AsyncClient
+from supabase import AClient as AsyncClient
 from leapfrogai_api.backend.rag.document_loader import load_file, split
 from leapfrogai_api.backend.rag.leapfrogai_embeddings import LeapfrogAIEmbeddings
 from leapfrogai_api.data.crud_file_bucket import CRUDFileBucket
@@ -347,7 +347,7 @@ async def aadd_documents(
                 metadata=document.metadata,
                 embedding=embedding,
             )
-            ids.append(response.data[0]["id"])
+            ids.append(response[0]["id"])
 
         return ids
 
@@ -450,5 +450,12 @@ async def _aadd_vector(
             "metadata": metadata,
             "embedding": embedding,
         }
-        response = await self.db.from_(self.table_name).insert(row).execute()
+        data, _count = await self.db.from_(self.table_name).insert(row).execute()
+
+        _, response = data
+
+        for item in response:
+            if "user_id" in item:
+                del item["user_id"]
+
         return response

src/leapfrogai_api/backend/rag/query.py

@@ -1,6 +1,6 @@
 """Service for querying the RAG model."""
 
-from supabase_py_async import AsyncClient
+from supabase import AClient as AsyncClient
 from leapfrogai_api.backend.rag.index import IndexingService
 from postgrest.base_request_builder import SingleAPIResponse
 

src/leapfrogai_api/backend/security/api_key.py

@@ -0,0 +1,71 @@
+"""API key pydantic model."""
+
+import secrets
+import hashlib
+from pydantic import BaseModel, field_validator, ValidationInfo
+
+KEY_PREFIX = "lfai"
+KEY_BYTES = 32
+CHECKSUM_LENGTH = 8
+
+
+class APIKey(BaseModel):
+    """API key model."""
+
+    prefix: str
+    unique_key: str
+    checksum: str
+
+    @classmethod
+    def generate(cls) -> "APIKey":
+        """Generate a new API key."""
+        unique_key: str = secrets.token_bytes(KEY_BYTES).hex()
+        checksum: str = cls._calculate_checksum(unique_key)
+        return cls(prefix=KEY_PREFIX, unique_key=unique_key, checksum=checksum)
+
+    @classmethod
+    def parse(cls, key_string: str) -> "APIKey":
+        """Parse a string representation of an API key."""
+        parts: list[str] = key_string.split("_")
+        if len(parts) != 3:
+            raise ValueError("Invalid API key format")
+        return cls(prefix=parts[0], unique_key=parts[1], checksum=parts[2])
+
+    @field_validator("prefix")
+    @classmethod
+    def validate_prefix(cls, prefix: str) -> str:
+        """Validate the key prefix."""
+        if prefix != KEY_PREFIX:
+            raise ValueError(f"Invalid prefix. Expected {KEY_PREFIX}")
+        return prefix
+
+    @field_validator("unique_key")
+    @classmethod
+    def validate_unique_key(cls, unique_key: str) -> str:
+        """Validate the unique key."""
+        if len(unique_key) != KEY_BYTES * 2:  # hex representation is twice as long
+            raise ValueError(
+                f"Invalid unique key length. Expected {KEY_BYTES * 2} characters"
+            )
+        return unique_key
+
+    @field_validator("checksum")
+    @classmethod
+    def validate_checksum(cls, checksum: str, info: ValidationInfo) -> str:
+        """Validate the checksum."""
+        if "unique_key" in info.data:
+            expected_checksum: str = cls._calculate_checksum(info.data["unique_key"])
+            if checksum != expected_checksum:
+                raise ValueError("Invalid checksum")
+        return checksum
+
+    def __str__(self) -> str:
+        return f"{self.prefix}_{self.unique_key}_{self.checksum}"
+
+    def __repr__(self) -> str:
+        return f"APIKey(prefix='{self.prefix}', unique_key='{self.unique_key}', checksum='{self.checksum}')"
+
+    @staticmethod
+    def _calculate_checksum(unique_key: str) -> str:
+        """Calculate a checksum for a unique key."""
+        return hashlib.sha256(unique_key.encode()).hexdigest()[:CHECKSUM_LENGTH]