planar 0.5.0__py3-none-any.whl

planar/db/db.py

@@ -0,0 +1,318 @@
+import asyncio
+import re
+from contextlib import asynccontextmanager
+from pathlib import Path
+from sqlite3 import LEGACY_TRANSACTION_CONTROL
+from typing import Any, Callable, Coroutine, cast
+
+from alembic import command
+from alembic.config import Config as AlembicConfig
+from pydantic import ConfigDict
+from sqlalchemy import Connection, MetaData, event, make_url, text
+from sqlalchemy.engine.url import URL
+from sqlalchemy.exc import DBAPIError
+from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
+from sqlalchemy.ext.compiler import compiles
+from sqlalchemy.orm import declared_attr
+from sqlalchemy.sql.expression import ClauseElement, Executable
+from sqlmodel import SQLModel
+from sqlmodel.ext.asyncio.session import AsyncSession
+
+import planar
+from planar.logging import get_logger
+from planar.modeling.orm.planar_base_entity import PLANAR_APPLICATION_METADATA
+from planar.utils import P, R, T, U, exponential_backoff_with_jitter
+
+
+def camel_to_snake(name):
+    name = re.sub("(.)([A-Z][a-z]+)", r"\1_\2", name)
+    return re.sub("([a-z0-9])([A-Z])", r"\1_\2", name).lower()
+
+
+PLANAR_SCHEMA = "planar"
+PLANAR_FRAMEWORK_METADATA = MetaData(schema=PLANAR_SCHEMA)
+logger = get_logger(__name__)
+
+
+class explain(Executable, ClauseElement):
+    inherit_cache = False
+
+    def __init__(self, stmt):
+        self.statement = stmt
+        self._inline = False
+
+
+@compiles(explain, "postgresql")
+def pg_explain(element, compiler, **kw):
+    text = "EXPLAIN ANALYZE "
+    text += compiler.process(element.statement, **kw)
+    return text
+
+
+class PlanarInternalBase(SQLModel, table=False):
+    """
+    Base model with common fields for all database tables.
+    Not a table itself - meant to be inherited by concrete model classes.
+
+    Usage conventions:
+    - Primary keys should be "id" and be UUID with default_factory=uuid4 when possible
+    - Use TimeStampMixin for auto-timestamp fields
+    - Field names should use snake_case consistently
+    - Table schema is set to 'planar' automatically
+    - Foreign keys should specify the full schema.table_name
+    """
+
+    @declared_attr.directive
+    def __tablename__(cls) -> str:  # type: ignore
+        return camel_to_snake(cls.__name__)
+
+    __abstract__ = True
+    # __table_args__ = {"schema": PLANAR_SCHEMA}
+    metadata = PLANAR_FRAMEWORK_METADATA
+    model_config = ConfigDict(validate_assignment=True)  # type: ignore
+
+
+class PlanarSession(AsyncSession):
+    def __init__(self, engine: AsyncEngine | None = None):
+        assert engine
+        self.engine = engine
+        self.dialect = engine.dialect
+        self.max_conflict_retries: int = 10
+        # dynamic import since planar.session depends on this
+        from planar.session import config_var
+
+        config = config_var.get(None)
+        if config is not None and config.app.max_db_conflict_retries:
+            self.max_conflict_retries = config.app.max_db_conflict_retries
+        super().__init__(engine, expire_on_commit=False)
+
+    async def set_serializable_isolation(self):
+        if self.dialect.name == "postgresql":
+            await self.exec(text("SET TRANSACTION ISOLATION LEVEL SERIALIZABLE"))  # type: ignore[arg-type]
+
+    @asynccontextmanager
+    async def begin_read(self):
+        """Context manager for read-only transactions.
+
+        This is useful when reading from the database since it ensures that if
+        a transaction has not started before the context, it will ensure no
+        transactions are open after the context.
+        """
+        in_transaction = self.in_transaction()
+        try:
+            yield
+            if not in_transaction:
+                await self.commit()
+        except Exception:
+            if not in_transaction:
+                await self.rollback()
+            raise
+
+    async def run_transaction(
+        self,
+        fn: Callable[P, Coroutine[T, U, R]],
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> R:
+        max_conflict_retries = self.max_conflict_retries
+
+        if self.in_transaction():
+            await self.commit()
+
+        remaining_retries = max_conflict_retries
+
+        while True:
+            try:
+                async with self.begin():
+                    await self.set_serializable_isolation()
+                    return await fn(*args, **kwargs)
+            except Exception as e:
+                if remaining_retries == 0:
+                    logger.exception("transaction failed after maximum retries")
+                    raise
+
+                if isinstance(e, DBAPIError) and "could not serialize access" in str(e):
+                    delay = exponential_backoff_with_jitter(
+                        max_conflict_retries - remaining_retries
+                    )
+                    await asyncio.sleep(delay)
+                    remaining_retries -= 1
+                    continue
+                logger.exception(
+                    "transaction failed due to unrecoverable error",
+                    remaining_retries=remaining_retries,
+                )
+                raise
+
+    async def explain(self, query: Executable, log_identifier: str) -> str:
+        if self.dialect.name != "postgresql" or not logger.isDebugEnabled():
+            return ""
+        # Reusing the current session will mess things up
+        # (implicit transaction starting, for example), so use
+        # a separate session to run the explain command
+        async with PlanarSession(self.engine) as session:
+            result = await session.exec(cast(Any, explain(query)))
+            query_plan = "\n".join([str(row[0]) for row in result])
+            compiled_sql = str(query)
+            logger.debug(
+                f"query_plan:{log_identifier}",
+                query_plan=query_plan,
+                compiled_sql=compiled_sql,
+            )
+            return query_plan
+
+
+def new_session(engine: AsyncEngine) -> PlanarSession:
+    return PlanarSession(engine)
+
+
+class DatabaseManager:
+    def __init__(
+        self,
+        db_url: str | URL,
+    ):
+        self.db_url = make_url(db_url) if isinstance(db_url, str) else db_url
+        self.engine: AsyncEngine | None = None
+
+    def _create_sqlite_engine(self, url: URL) -> AsyncEngine:
+        # in practice this high timeout is only use
+        timeout = int(str(url.query.get("timeout", 10)))
+
+        engine = create_async_engine(
+            url,
+            connect_args=dict(
+                timeout=timeout,
+                isolation_level=None,
+                # If autocommit is not LEGACY_TRANSACTION_CONTROL, isolation_level
+                # is ignored, so we set here explicitly to make the intention clear,
+                # even though it is the default value.
+                autocommit=LEGACY_TRANSACTION_CONTROL,
+            ),
+            # SQLite doesn't support schemas, so we need to translate the planar schema
+            # name to None.
+            execution_options={"schema_translate_map": {"planar": None}},
+        )
+
+        def do_begin(conn: Connection):
+            conn.exec_driver_sql("BEGIN IMMEDIATE")
+
+        event.listen(engine.sync_engine, "begin", do_begin)
+
+        return engine
+
+    def _create_postgresql_engine(self, url: URL) -> AsyncEngine:
+        engine = create_async_engine(url)
+
+        return engine
+
+    def connect(self):
+        """Creates and initializes the database engine."""
+        if self.engine:
+            logger.warning("database engine already initialized")
+            return
+
+        db_backend = self.db_url.get_backend_name()
+
+        match db_backend:
+            case "sqlite":
+                logger.info(
+                    "connecting to database", db_backend=db_backend, db_url=self.db_url
+                )
+                self.engine = self._create_sqlite_engine(self.db_url)
+            case "postgresql":
+                logger.info("connecting to database", db_backend=db_backend)
+                self.engine = self._create_postgresql_engine(self.db_url)
+            case _:
+                raise NotImplementedError(
+                    f'Unsupported database backend "{db_backend}"'
+                )
+
+    async def disconnect(self):
+        """Disposes of the database engine."""
+        if self.engine:
+            logger.info("disconnecting database engine")
+            await self.engine.dispose()
+            self.engine = None
+        else:
+            logger.warning("attempted to disconnect an uninitialized engine")
+
+    def get_engine(self) -> AsyncEngine:
+        """Returns the initialized AsyncEngine."""
+        if not self.engine:
+            raise RuntimeError("Database engine not initialized. Call connect() first.")
+        return self.engine
+
+    def get_session(self) -> PlanarSession:
+        """Returns a new PlanarSession."""
+        if not self.engine:
+            raise RuntimeError("Database engine not initialized. Call connect() first.")
+        return PlanarSession(self.engine)
+
+    async def _run_system_migrations(self):
+        logger.info("running planar system migrations")
+
+        module_path = Path(planar.__file__).parent
+        script_location = str(module_path / "db" / "alembic")
+
+        alembic_cfg = AlembicConfig()
+        alembic_cfg.set_main_option("script_location", script_location)
+
+        if not self.engine:
+            raise RuntimeError("Database engine not initialized. Call connect() first.")
+        try:
+            async with self.engine.begin() as conn:
+                # Pass the *synchronous* connection produced by `run_sync` to Alembic.
+
+                def _upgrade(sync_conn):
+                    """Run Alembic upgrade using the given synchronous connection."""
+
+                    # Inject the sync SQLAlchemy Connection so that planar/db/alembic/env.py
+                    # recognises we're running in programmatic (runtime) mode instead of
+                    # development mode. This prevents it from trying to create a new engine
+                    # via `engine_from_config`, which expects a URL in the Alembic config.
+                    alembic_cfg.attributes["connection"] = sync_conn
+
+                    # Execute migrations up to the latest revision.
+                    command.upgrade(alembic_cfg, "head")
+
+                # Execute the upgrade inside the green-thread aware sync context.
+                await conn.run_sync(_upgrade)
+            logger.info("planar system migrations completed successfully")
+        except Exception:
+            logger.exception("planar system migration failed")
+            raise
+
+    async def _setup_database(self):
+        if not self.engine:
+            raise RuntimeError("Database engine not initialized. Call connect() first.")
+
+        async with self.engine.begin() as conn:
+            if "sqlite" in self.db_url.drivername:
+                await conn.execute(text("PRAGMA foreign_keys=ON"))
+            else:
+                # Ensure planar schema exists
+                await conn.execute(text(f"CREATE SCHEMA IF NOT EXISTS {PLANAR_SCHEMA}"))
+
+    async def migrate(self, use_alembic: bool):
+        """
+        Runs database migrations.
+        By default, uses SQLModel.metadata.create_all.
+        Set use_alembic=True to use Alembic (requires Alembic setup).
+        """
+        if not self.engine:
+            raise RuntimeError("Database engine not initialized. Call connect() first.")
+
+        logger.info("starting database migration")
+        if use_alembic:
+            logger.info("using alembic for migrations")
+            await self._setup_database()
+            await self._run_system_migrations()
+            # For now user migrations are not supported, so we fall back to SQLModel.metadata.create_all
+            async with self.engine.begin() as conn:
+                await conn.run_sync(PLANAR_APPLICATION_METADATA.create_all)
+
+        else:
+            async with self.engine.begin() as conn:
+                await self._setup_database()
+                await conn.run_sync(PLANAR_FRAMEWORK_METADATA.create_all)
+                await conn.run_sync(PLANAR_APPLICATION_METADATA.create_all)

planar/files/__init__.py

@@ -0,0 +1,2 @@
+from .models import PlanarFile, PlanarFileMetadata  # noqa: F401
+from .storage.context import get_storage  # noqa: F401

planar/files/models.py

@@ -0,0 +1,162 @@
+import mimetypes
+import os
+from pathlib import Path
+from typing import AsyncGenerator, Union
+from uuid import UUID, uuid4
+
+import aiofiles
+from pydantic import BaseModel
+from sqlmodel import Field
+
+from planar.db import PlanarInternalBase
+from planar.files.storage.context import get_storage
+from planar.logging import get_logger
+from planar.modeling.mixins import TimestampMixin
+from planar.session import get_session
+
+logger = get_logger(__name__)
+
+
+class PlanarFile(BaseModel):
+    id: UUID
+    filename: str
+    content_type: str
+    size: int
+
+    async def get_metadata(self) -> "PlanarFileMetadata":
+        """
+        Retrieves the metadata for this file from the database.
+        """
+        logger.debug("getting metadata for file", file_id=self.id)
+        session = get_session()
+        async with session.begin_read():
+            result = await session.get(PlanarFileMetadata, self.id)
+        if result is None:
+            logger.warning("file metadata not found in database", file_id=self.id)
+            raise ValueError(f"File with ID {self.id} not found in the database.")
+        return result
+
+    async def get_content(self) -> bytes:
+        """
+        Retrieves the content of this file from the storage backend.
+        """
+        logger.debug("getting content for file", file_id=self.id)
+        storage = get_storage()
+        metadata = await self.get_metadata()
+        data, _ = await storage.get_bytes(metadata.storage_ref)
+        return data
+
+    @staticmethod
+    async def upload(
+        content: Union[bytes, AsyncGenerator[bytes, None], Path, str],
+        filename: str,
+        content_type: str | None = None,
+        size: int | None = None,
+    ) -> "PlanarFile":
+        """
+        Uploads file content to storage and creates its metadata record.
+
+        Args:
+            content: File content as bytes, an async iterator, or a file path (str or Path).
+            filename: The desired filename for storage and metadata.
+            content_type: The MIME type of the file. If None, it's inferred from the filename
+                          for paths or defaults to 'application/octet-stream'.
+            size: The size of the file in bytes. If None, it's calculated for bytes/paths
+                  or defaults to -1 for streams.
+
+        Returns:
+            The created PlanarFile object with metadata.
+
+        Raises:
+            FileNotFoundError: If content is a path and the file doesn't exist.
+            TypeError: If the content type is not supported.
+        """
+        logger.debug(
+            "uploading file",
+            filename=filename,
+            content_type=content_type,
+            size=size,
+        )
+        storage = get_storage()
+        session = get_session()
+
+        storage_ref: str
+        actual_size: int = -1
+        final_content_type: str = content_type or "application/octet-stream"
+
+        if isinstance(content, (str, Path)):
+            file_path = Path(content)
+            logger.debug("uploading from path", path=file_path)
+            if not file_path.is_file():
+                logger.warning("file not found at path for upload", path=file_path)
+                raise FileNotFoundError(f"File not found at path: {file_path}")
+
+            actual_size = size if size is not None else os.path.getsize(file_path)
+
+            if content_type is None:
+                guessed_type, _ = mimetypes.guess_type(filename)
+                final_content_type = guessed_type or "application/octet-stream"
+
+            async def file_stream():
+                async with aiofiles.open(file_path, "rb") as afp:
+                    chunk_size = 65536  # 64KB chunk size
+                    while chunk := await afp.read(chunk_size):
+                        yield chunk
+
+            storage_ref = await storage.put(
+                stream=file_stream(), mime_type=final_content_type
+            )
+
+        elif isinstance(content, bytes):
+            logger.debug("uploading from bytes")
+            actual_size = size if size is not None else len(content)
+            # Keep provided content_type or default
+            final_content_type = content_type or "application/octet-stream"
+            storage_ref = await storage.put_bytes(content, mime_type=final_content_type)
+
+        elif isinstance(content, AsyncGenerator):  # Check for async iterator
+            logger.debug("uploading from async generator stream")
+            actual_size = size if size is not None else -1  # Size required or unknown
+            # Keep provided content_type or default
+            final_content_type = content_type or "application/octet-stream"
+            storage_ref = await storage.put(
+                stream=content, mime_type=final_content_type
+            )
+        else:
+            logger.warning(
+                "unsupported content type for upload", content_type=type(content)
+            )
+            raise TypeError(
+                "Unsupported content type. Must be bytes, AsyncGenerator, str path, or Path object."
+            )
+
+        # Create the metadata record
+        planar_file_metadata = PlanarFileMetadata(
+            filename=filename,
+            content_type=final_content_type,
+            size=actual_size,
+            storage_ref=storage_ref,
+        )
+        session.add(planar_file_metadata)
+        await session.commit()
+        await session.refresh(planar_file_metadata)
+        logger.info(
+            "file uploaded and metadata created",
+            id=planar_file_metadata.id,
+            filename=filename,
+            storage_ref=storage_ref,
+        )
+
+        # We return the metadata instance which also satisfies the PlanarFile structure
+        return planar_file_metadata
+
+
+class PlanarFileMetadata(PlanarFile, TimestampMixin, PlanarInternalBase, table=True):
+    """
+    Database model storing the authoritative mapping between a PlanarFile.file_id
+    and its storage details. Acts as the single, central file manifest.
+    """
+
+    id: UUID = Field(default_factory=uuid4, primary_key=True)
+    # storage_ref is a storage backend specifid identifier for the file
+    storage_ref: str = Field(index=True)

planar/files/storage/base.py

@@ -0,0 +1,61 @@
+import io
+from abc import ABC, abstractmethod
+from typing import AsyncGenerator
+
+
+class Storage(ABC):
+    @abstractmethod
+    async def put(
+        self, stream: AsyncGenerator[bytes, None], mime_type: str | None = None
+    ) -> str:
+        """Store a stream and its mime type, returning a storage ref."""
+        ...
+
+    @abstractmethod
+    async def get(self, ref: str) -> tuple[AsyncGenerator[bytes, None], str | None]:
+        """Get a stream and its mime type from a storage ref."""
+        ...
+
+    @abstractmethod
+    async def delete(self, ref: str) -> None:
+        """Delete the object associated with the storage ref."""
+        ...
+
+    @abstractmethod
+    async def external_url(self, ref: str) -> str | None:
+        """If available, return an external URL to read the file."""
+        ...
+
+    async def put_bytes(self, data: bytes, mime_type: str | None = None) -> str:
+        """Store bytes and optional mime type, returning a storage ref."""
+
+        async def _stream():
+            yield data
+
+        return await self.put(_stream(), mime_type=mime_type)
+
+    async def get_bytes(self, ref: str) -> tuple[bytes, str | None]:
+        """Get bytes and mime type from a storage ref."""
+        buffer = io.BytesIO()
+        stream, mime_type = await self.get(ref)
+        async for chunk in stream:
+            buffer.write(chunk)
+        return buffer.getvalue(), mime_type
+
+    async def put_string(
+        self, data: str, encoding: str = "utf-8", mime_type: str | None = None
+    ) -> str:
+        """Store a string and optional mime type, returning a storage ref."""
+        # Ensure mime_type includes encoding if not already specified
+        final_mime_type = mime_type
+        if mime_type and "charset=" not in mime_type and mime_type.startswith("text/"):
+            final_mime_type = f"{mime_type}; charset={encoding}"
+        return await self.put_bytes(data.encode(encoding), mime_type=final_mime_type)
+
+    async def get_string(
+        self, ref: str, encoding: str = "utf-8"
+    ) -> tuple[str, str | None]:
+        """Get a string and mime type from a storage ref."""
+        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

planar/files/storage/config.py

@@ -0,0 +1,44 @@
+from typing import Annotated, Literal, Optional
+
+from pydantic import BaseModel, Field
+
+from .local_directory import LocalDirectoryStorage
+from .s3 import S3Storage
+
+
+class LocalDirectoryConfig(BaseModel):
+    backend: Literal["localdir"]
+    directory: str
+
+
+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
+    presigned_url_ttl: int = 3600
+
+
+StorageConfig = Annotated[
+    LocalDirectoryConfig | S3Config,
+    Field(discriminator="backend"),
+]
+
+
+def create_from_config(config: StorageConfig) -> LocalDirectoryStorage | S3Storage:
+    """Creates a storage instance from the given configuration."""
+    if config.backend == "localdir":
+        return LocalDirectoryStorage(config.directory)
+    elif config.backend == "s3":
+        return S3Storage(
+            bucket_name=config.bucket_name,
+            region=config.region,
+            access_key_id=config.access_key,
+            secret_access_key=config.secret_key,
+            endpoint_url=config.endpoint_url,
+            presigned_url_ttl=config.presigned_url_ttl,
+        )
+    else:
+        raise ValueError(f"Unsupported backend: {config.backend}")

planar/files/storage/context.py

@@ -0,0 +1,15 @@
+from contextvars import ContextVar
+
+from planar.files.storage.base import Storage
+
+storage_var: ContextVar[Storage] = ContextVar("storage")
+
+
+def get_storage() -> Storage:
+    """Get the current storage context."""
+    return storage_var.get()
+
+
+def set_storage(storage: Storage) -> None:
+    """Set the current storage context."""
+    storage_var.set(storage)