planar 0.5.0__py3-none-any.whl → 0.7.0__py3-none-any.whl

planar/config.py

@@ -5,13 +5,14 @@ import os
 import sys
 from enum import Enum
 from pathlib import Path
-from typing import Annotated, Any, Dict, Literal, Optional
+from typing import Annotated, Any, Dict, Literal
 
 import boto3
 import yaml
 from dotenv import load_dotenv
 from pydantic import (
     BaseModel,
+    ConfigDict,
     Field,
     HttpUrl,
     SecretStr,
@@ -45,8 +46,8 @@ class LogLevel(str, Enum):
 
 class LoggerConfig(BaseModel):
     level: LogLevel = LogLevel.INFO
-    propagate: Optional[bool] = False
-    file: Optional[str] = None
+    propagate: bool | None = False
+    file: str | None = None
 
 
 class SQLiteConfig(BaseModel):
@@ -66,11 +67,11 @@ class PostgreSQLConfig(BaseModel):
     driver: Literal["postgresql", "postgresql+asyncpg"] = (
         "postgresql+asyncpg"  # Allow async PostgreSQL
     )
-    host: Optional[str] = None
-    port: Optional[int] = None
-    user: Optional[str] = None
-    password: Optional[str] = None
-    db: Optional[str]
+    host: str | None = None
+    port: int | None = None
+    user: str | None = None
+    password: str | None = None
+    db: str | None
 
     def connection_url(self) -> URL:
         driver = self.driver
@@ -92,15 +93,15 @@ class OpenAIConfig(BaseModel):
     """Configuration for OpenAI provider."""
 
     api_key: SecretStr
-    base_url: Optional[str] = None
-    organization: Optional[str] = None
+    base_url: str | None = None
+    organization: str | None = None
 
 
 class AnthropicConfig(BaseModel):
     """Configuration for Anthropic provider."""
 
     api_key: SecretStr
-    base_url: Optional[str] = None
+    base_url: str | None = None
 
 
 class GeminiConfig(BaseModel):
@@ -112,9 +113,9 @@ class GeminiConfig(BaseModel):
 class AIProvidersConfig(BaseModel):
     """Configuration for AI providers."""
 
-    openai: Optional[OpenAIConfig] = None
-    anthropic: Optional[AnthropicConfig] = None
-    gemini: Optional[GeminiConfig] = None
+    openai: OpenAIConfig | None = None
+    anthropic: AnthropicConfig | None = None
+    gemini: GeminiConfig | None = None
 
 
 DatabaseConfig = Annotated[
@@ -124,7 +125,7 @@ DatabaseConfig = Annotated[
 
 class AppConfig(BaseModel):
     db_connection: str
-    max_db_conflict_retries: Optional[int] = None
+    max_db_conflict_retries: int | None = None
 
 
 def default_storage_config() -> StorageConfig:
@@ -162,31 +163,27 @@ PROD_CORS_CONFIG = CorsConfig(
 
 
 class JWTConfig(BaseModel):
-    enabled: bool = False
     client_id: str | None = None
     org_id: str | None = None
     additional_exclusion_paths: list[str] | None = Field(default_factory=list)
 
     @model_validator(mode="after")
     def validate_client_id(cls, instance):
-        if instance.enabled and not instance.client_id:
-            raise ValueError("client_id is required when JWT is enabled")
-        if instance.client_id and not instance.enabled:
-            raise ValueError(
-                "You cannot specify a client_id without enabling JWT - did you mean to set enabled=True?"
-            )
+        if not instance.client_id or not instance.org_id:
+            raise ValueError("Both client_id and org_id required to enable JWT")
         return instance
 
 
-JWT_DISABLED_CONFIG = JWTConfig(enabled=False)
+# Coplane ORG JWT config
 JWT_COPLANE_CONFIG = JWTConfig(
-    enabled=True, client_id="client_01JSJHJP9Q8GZDK5Y856FEHTB0", org_id=None
+    client_id="client_01JSJHJPKG09TMSK6NHJP0S180",
+    org_id="org_01JY4QP57Y7H4EQ7HT3BGN7TNK",
 )
 
 
 class OtelConfig(BaseModel):
     collector_endpoint: HttpUrl
-    resource_attributes: Optional[dict[str, str]] = None
+    resource_attributes: dict[str, str] | None = None
 
 
 def install_otel_provider(otel_config: OtelConfig):
@@ -206,19 +203,31 @@ class AuthzConfig(BaseModel):
     policy_file: str | None = None
 
 
+class ServiceTokenConfig(BaseModel):
+    token: str | None = Field(None, min_length=1)
+
+
+class SecurityConfig(BaseModel):
+    cors: CorsConfig = PROD_CORS_CONFIG
+    jwt: JWTConfig | None = None
+    service_token: ServiceTokenConfig | None = None
+    authz: AuthzConfig | None = None
+
+
 class PlanarConfig(BaseModel):
     db_connections: Dict[str, DatabaseConfig | str]
     app: AppConfig
-    ai_providers: Optional[AIProvidersConfig] = None
-    storage: Optional[StorageConfig] = default_storage_config()
+    ai_providers: AIProvidersConfig | None = None
+    storage: StorageConfig | None = default_storage_config()
     sse_hub: str | bool = False
-    cors: CorsConfig = PROD_CORS_CONFIG
     environment: Environment = Environment.DEV
-    jwt: JWTConfig | None = None
-    logging: Optional[dict[str, LoggerConfig]] = None
+    security: SecurityConfig = SecurityConfig()
+    logging: dict[str, LoggerConfig] | None = None
     use_alembic: bool | None = True
-    otel: Optional[OtelConfig] = None
-    authz: AuthzConfig | None = None
+    otel: OtelConfig | None = None
+
+    # forbid extra keys in the config to prevent accidental misconfiguration
+    model_config = ConfigDict(extra="forbid")
 
     @model_validator(mode="after")
     def validate_db_connection_reference(cls, instance):
@@ -467,14 +476,14 @@ def load_environment_aware_config[ConfigClass]() -> PlanarConfig:
 
     if env == "dev":
         base_config = sqlite_config(db_path="planar_dev.db")
-        base_config.cors = LOCAL_CORS_CONFIG
+        base_config.security = SecurityConfig(cors=LOCAL_CORS_CONFIG)
         base_config.environment = Environment.DEV
-        base_config.jwt = JWT_DISABLED_CONFIG
     else:
         base_config = sqlite_config(db_path="planar.db")
-        base_config.cors = PROD_CORS_CONFIG
         base_config.environment = Environment.PROD
-        base_config.jwt = JWT_COPLANE_CONFIG
+        base_config.security = SecurityConfig(
+            cors=PROD_CORS_CONFIG, jwt=JWT_COPLANE_CONFIG
+        )
 
     # Convert base config to dict for merging
     # Use by_alias=False to work with Python field names before validation

planar/db/db.py

@@ -176,7 +176,8 @@ class DatabaseManager:
 
     def _create_sqlite_engine(self, url: URL) -> AsyncEngine:
         # in practice this high timeout is only use
-        timeout = int(str(url.query.get("timeout", 10)))
+        timeout = int(str(url.query.get("timeout", 60)))
+        logger.info("Setting up SQLite engine with timeout", timeout=timeout)
 
         engine = create_async_engine(
             url,

planar/files/storage/azure_blob.py

@@ -0,0 +1,343 @@
+import uuid
+from datetime import UTC, datetime, timedelta
+from enum import Enum
+from typing import AsyncGenerator, override
+from urllib.parse import urlparse
+
+from azure.core.exceptions import ResourceNotFoundError
+from azure.storage.blob import BlobSasPermissions, generate_blob_sas
+from azure.storage.blob.aio import BlobServiceClient
+
+from planar.logging import get_logger
+
+from .base import Storage
+
+logger = get_logger(__name__)
+
+
+class AzureAuthMethod(Enum):
+    CONNECTION_STRING = "connection_string"
+    ACCOUNT_KEY = "account_key"
+    AZURE_AD = "azure_ad"
+
+
+class AzureBlobStorage(Storage):
+    """Stores files and mime types in Azure Blob Storage using the async SDK."""
+
+    def __init__(
+        self,
+        container_name: str,
+        connection_string: str | None = None,
+        account_url: str | None = None,
+        use_azure_ad: bool = False,
+        account_key: str | None = None,
+        sas_ttl: int = 3600,
+    ):
+        """
+        Initializes AzureBlobStorage.
+
+        Args:
+            container_name: The name of the Azure Storage container.
+            connection_string: Full connection string (includes all credentials).
+            account_url: Storage account URL (e.g., 'https://<account>.blob.core.windows.net').
+            use_azure_ad: Whether to use DefaultAzureCredential for Azure AD auth.
+            account_key: Storage account key (used with account_url).
+            sas_ttl: Time in seconds for which SAS URLs are valid.
+        """
+        # Import Azure dependencies when the class is instantiated
+        try:
+            from azure.storage.blob.aio import BlobServiceClient
+        except ImportError as e:
+            raise ImportError(
+                "Azure storage dependencies are not installed. "
+                "Install with: pip install planar[azure]"
+            ) from e
+
+        self.container_name = container_name
+        self.sas_ttl = sas_ttl
+        self.client: "BlobServiceClient"
+
+        self.auth_method: AzureAuthMethod
+        self._account_name: str | None = None
+        self._account_key: str | None = None
+
+        from azure.storage.blob._shared.policies_async import ExponentialRetry
+
+        client_kwargs = {
+            "connection_timeout": 10,
+            "read_timeout": 40,
+            "retry_policy": ExponentialRetry(
+                retry_total=2,
+            ),
+        }
+
+        # Initialize BlobServiceClient based on auth method
+        if connection_string:
+            self.client = BlobServiceClient.from_connection_string(
+                connection_string,
+                **client_kwargs,
+            )
+            self.auth_method = AzureAuthMethod.CONNECTION_STRING
+            # Extract account name and key from the connection string for SAS signing
+            self._account_name = self._extract_account_name_from_connection_string(
+                connection_string
+            )
+            self._account_key = self._extract_account_key_from_connection_string(
+                connection_string
+            )
+
+        elif use_azure_ad:
+            if not account_url:
+                raise ValueError(
+                    "account_url is required when using Azure AD authentication"
+                )
+            from azure.identity.aio import DefaultAzureCredential
+
+            credential = DefaultAzureCredential()
+            self.client = BlobServiceClient(
+                account_url=account_url, credential=credential, **client_kwargs
+            )
+            self.auth_method = AzureAuthMethod.AZURE_AD
+            self._account_key = None
+            self._account_name = self._extract_account_name_from_account_url(
+                account_url
+            )
+
+        elif account_key:
+            if not account_url:
+                raise ValueError(
+                    "account_url is required when using account key authentication"
+                )
+            self.client = BlobServiceClient(
+                account_url=account_url, credential=account_key, **client_kwargs
+            )
+            self.auth_method = AzureAuthMethod.ACCOUNT_KEY
+            self._account_key = account_key
+            # Extract account name from URL for SAS generation
+            self._account_name = self._extract_account_name_from_account_url(
+                account_url
+            )
+
+        else:
+            raise ValueError(
+                "Must provide either connection_string, use_azure_ad=True, or account_key"
+            )
+
+    async def __aenter__(self):
+        """Enter async context manager for proper cleanup in tests."""
+        await self.client.__aenter__()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Exit async context manager and cleanup resources."""
+        await self.client.__aexit__(exc_type, exc_val, exc_tb)
+
+    @override
+    async def close(self):
+        """Explicitly close the client. Only needed if not using as context manager."""
+        await self.client.close()
+
+    async def put(
+        self, stream: AsyncGenerator[bytes, None], mime_type: str | None = None
+    ) -> str:
+        """
+        Stores a stream and optional mime type to Azure Blob Storage.
+
+        The storage reference returned is a unique UUID.
+        The mime_type is stored as the blob's ContentType.
+        """
+        ref = str(uuid.uuid4())
+
+        content_settings = None
+        if mime_type:
+            from azure.storage.blob import ContentSettings
+
+            content_settings = ContentSettings(content_type=mime_type)
+
+        try:
+            container_client = self.client.get_container_client(self.container_name)
+            blob_client = container_client.get_blob_client(ref)
+
+            await blob_client.upload_blob(
+                stream,
+                content_settings=content_settings,
+                overwrite=True,
+            )
+            return ref
+
+        except Exception as e:
+            logger.exception(
+                "failed azure blob upload",
+                ref=ref,
+                container_name=self.container_name,
+            )
+            raise IOError(f"Failed to upload to Azure blob {ref}. Error: {e}") from e
+
+    async def get(self, ref: str) -> tuple[AsyncGenerator[bytes, None], str | None]:
+        """
+        Retrieves a stream and its mime type from Azure Blob Storage.
+        """
+        try:
+            container_client = self.client.get_container_client(self.container_name)
+            blob_client = container_client.get_blob_client(ref)
+
+            # Get blob properties for content type
+            properties = await blob_client.get_blob_properties()
+            mime_type = (
+                properties.content_settings.content_type
+                if properties.content_settings
+                else None
+            )
+
+            async def _stream_wrapper():
+                download_stream = await blob_client.download_blob()
+                async for chunk in download_stream.chunks():
+                    yield chunk
+
+            return _stream_wrapper(), mime_type
+        except ResourceNotFoundError as e:
+            logger.warning(
+                "azure blob not found",
+                ref=ref,
+                container_name=self.container_name,
+                error=e,
+            )
+            raise FileNotFoundError(f"Azure blob not found: {ref}") from e
+        except Exception as e:
+            logger.exception(
+                "failed azure blob download",
+                ref=ref,
+                container_name=self.container_name,
+            )
+            raise IOError(
+                f"Failed to download from Azure blob {ref}. Error: {e}"
+            ) from e
+
+    async def delete(self, ref: str) -> None:
+        """
+        Deletes a blob from Azure Storage.
+        Does not raise an error if the blob does not exist.
+        """
+        try:
+            container_client = self.client.get_container_client(self.container_name)
+            blob_client = container_client.get_blob_client(ref)
+
+            await blob_client.delete_blob(delete_snapshots="include")
+
+        except ResourceNotFoundError:
+            logger.debug(
+                "azure blob not found, not raising error",
+                ref=ref,
+                container_name=self.container_name,
+            )
+        except Exception as e:
+            logger.exception(
+                "failed azure blob delete",
+                ref=ref,
+                container_name=self.container_name,
+            )
+            raise IOError(f"Failed to delete Azure blob {ref}. Error: {e}") from e
+
+    async def external_url(self, ref: str) -> str | None:
+        """
+        Returns a SAS URL to access the Azure blob if we have the capability.
+
+        Supports SAS generation for:
+        - Account Key (Account SAS signed with account key)
+        - Connection String (Account SAS signed with extracted account key)
+        - Azure AD (User Delegation SAS signed with a User Delegation Key)
+        """
+
+        if not self._account_name:
+            logger.debug(
+                "cannot generate sas url without account name",
+                ref=ref,
+                has_account_name=bool(self._account_name),
+            )
+            return None
+
+        expiry_time = datetime.now(UTC) + timedelta(seconds=self.sas_ttl)
+
+        if self.auth_method.name in ("ACCOUNT_KEY", "CONNECTION_STRING"):
+            if not self._account_key:
+                logger.debug(
+                    "cannot generate account-key SAS without account key",
+                    ref=ref,
+                    has_account_key=bool(self._account_key),
+                )
+                return None
+
+            sas_token = generate_blob_sas(
+                account_name=self._account_name,
+                container_name=self.container_name,
+                blob_name=ref,
+                account_key=self._account_key,
+                permission=BlobSasPermissions(read=True),
+                expiry=expiry_time,
+            )
+
+        elif self.auth_method.name == "AZURE_AD":
+            # Generate a User Delegation SAS signed with a user delegation key
+            start_time = datetime.utcnow()
+            user_delegation_key = await self.client.get_user_delegation_key(
+                key_start_time=start_time, key_expiry_time=expiry_time
+            )
+            sas_token = generate_blob_sas(
+                account_name=self._account_name,
+                container_name=self.container_name,
+                blob_name=ref,
+                user_delegation_key=user_delegation_key,
+                permission=BlobSasPermissions(read=True),
+                expiry=expiry_time,
+            )
+        else:
+            return None
+
+        blob_url = f"{self.client.url}{self.container_name}/{ref}"
+        return f"{blob_url}?{sas_token}"
+
+    @staticmethod
+    def _extract_account_name_from_connection_string(
+        connection_string: str,
+    ) -> str | None:
+        try:
+            # Split on ';' and build a dict of key/value pairs
+            parts = dict(
+                part.split("=", 1)
+                for part in connection_string.split(";")
+                if "=" in part
+            )
+            account_name = parts.get("AccountName")
+            return account_name
+        except Exception:
+            return None
+
+    @staticmethod
+    def _extract_account_key_from_connection_string(
+        connection_string: str,
+    ) -> str | None:
+        try:
+            parts = dict(
+                part.split("=", 1)
+                for part in connection_string.split(";")
+                if "=" in part
+            )
+            return parts.get("AccountKey")
+        except Exception:
+            return None
+
+    @staticmethod
+    def _extract_account_name_from_account_url(account_url: str) -> str | None:
+        try:
+            parsed = urlparse(account_url)
+            host = parsed.hostname or ""
+            # Standard Azure: https://{account}.blob.core.windows.net
+            if "." in host and not host.startswith("127.") and host != "localhost":
+                return host.split(".")[0]
+            # Azurite style: http://127.0.0.1:10000/{account}
+            path = parsed.path.strip("/")
+            if path:
+                return path.split("/")[0]
+            return None
+        except Exception:
+            return None

planar/files/storage/base.py

@@ -59,3 +59,10 @@ class Storage(ABC):
         data_bytes, mime_type = await self.get_bytes(ref)
         # TODO: Potentially use encoding from mime_type if available?
         return data_bytes.decode(encoding), mime_type
+
+    async def close(self) -> None:
+        """
+        Optional cleanup method for storage implementations.
+        Override this if your storage backend needs explicit cleanup.
+        """
+        pass

planar/files/storage/config.py

@@ -1,10 +1,15 @@
-from typing import Annotated, Literal, Optional
+from __future__ import annotations
 
-from pydantic import BaseModel, Field
+from typing import TYPE_CHECKING, Annotated, Literal
+
+from pydantic import BaseModel, Field, model_validator
 
 from .local_directory import LocalDirectoryStorage
 from .s3 import S3Storage
 
+if TYPE_CHECKING:
+    from .azure_blob import AzureBlobStorage
+
 
 class LocalDirectoryConfig(BaseModel):
     backend: Literal["localdir"]
@@ -15,19 +20,66 @@ class S3Config(BaseModel):
     backend: Literal["s3"]
     bucket_name: str
     region: str
-    access_key: Optional[str] = None
-    secret_key: Optional[str] = None
-    endpoint_url: Optional[str] = None
+    access_key: str | None = None
+    secret_key: str | None = None
+    endpoint_url: str | None = None
     presigned_url_ttl: int = 3600
 
 
+class AzureBlobConfig(BaseModel):
+    backend: Literal["azure_blob"]
+    container_name: str
+
+    # Authentication options (mutually exclusive)
+    connection_string: str | None = None  # Full connection string
+    account_url: str | None = None  # Storage account URL
+    use_azure_ad: bool | None = None  # Use DefaultAzureCredential
+    account_key: str | None = None  # Storage account key
+
+    # Common settings
+    sas_ttl: int = 3600  # SAS URL expiry time in seconds
+
+    @model_validator(mode="after")
+    def validate_auth_config(self):
+        """Ensure exactly one valid authentication configuration."""
+
+        # Check if connection_string is provided
+        if self.connection_string:
+            # Connection string is self-contained
+            if self.account_url or self.use_azure_ad or self.account_key:
+                raise ValueError(
+                    "When using connection_string, don't provide account_url, use_azure_ad, or account_key"
+                )
+            return self
+
+        # If no connection string, must have account_url
+        if not self.account_url:
+            raise ValueError("Either connection_string or account_url must be provided")
+
+        # With account_url, must have exactly one credential type
+        credential_methods = [
+            self.use_azure_ad is True,
+            self.account_key is not None,
+        ]
+
+        if sum(credential_methods) != 1:
+            raise ValueError(
+                "When using account_url, exactly one credential method must be specified: "
+                "use_azure_ad=true or account_key"
+            )
+
+        return self
+
+
 StorageConfig = Annotated[
-    LocalDirectoryConfig | S3Config,
+    LocalDirectoryConfig | S3Config | AzureBlobConfig,
     Field(discriminator="backend"),
 ]
 
 
-def create_from_config(config: StorageConfig) -> LocalDirectoryStorage | S3Storage:
+def create_from_config(
+    config: StorageConfig,
+) -> LocalDirectoryStorage | S3Storage | AzureBlobStorage:
     """Creates a storage instance from the given configuration."""
     if config.backend == "localdir":
         return LocalDirectoryStorage(config.directory)
@@ -40,5 +92,16 @@ def create_from_config(config: StorageConfig) -> LocalDirectoryStorage | S3Stora
             endpoint_url=config.endpoint_url,
             presigned_url_ttl=config.presigned_url_ttl,
         )
+    elif config.backend == "azure_blob":
+        from .azure_blob import AzureBlobStorage
+
+        return AzureBlobStorage(
+            container_name=config.container_name,
+            connection_string=config.connection_string,
+            account_url=config.account_url,
+            use_azure_ad=config.use_azure_ad or False,
+            account_key=config.account_key,
+            sas_ttl=config.sas_ttl,
+        )
     else:
         raise ValueError(f"Unsupported backend: {config.backend}")

planar/files/storage/s3.py

@@ -1,7 +1,7 @@
 import asyncio
 import io
 import uuid
-from typing import Any, AsyncGenerator, Dict, Optional, Tuple
+from typing import Any, AsyncGenerator, Dict, Tuple
 
 import boto3
 from botocore.client import Config as BotoConfig
@@ -21,11 +21,11 @@ class S3Storage(Storage):
         self,
         bucket_name: str,
         region: str,
-        endpoint_url: Optional[str] = None,
-        access_key_id: Optional[str] = None,
-        secret_access_key: Optional[str] = None,
-        session_token: Optional[str] = None,  # For temporary credentials
-        boto_config: Optional[Dict[str, Any]] = None,  # Additional boto3 client config
+        endpoint_url: str | None = None,
+        access_key_id: str | None = None,
+        secret_access_key: str | None = None,
+        session_token: str | None = None,  # For temporary credentials
+        boto_config: Dict[str, Any] | None = None,  # Additional boto3 client config
         presigned_url_ttl: int = 3600,
     ):
         """