Add pydantic models

pyproject.toml

@@ -8,7 +8,8 @@ dependencies = [
   "dockerspawner~=12.1",
   "jupyter_client>=6.1,<8",
   "httpx",
-  "sqlalchemy>=2"
+  "sqlalchemy>=2",
+  "pydantic>=2,<3"
 ]
 dynamic = ["version"]
 license = {file = "LICENSE"}

tljh_repo2docker/alembic/versions/ac1b4e7e52f3_first_migration.py

@@ -15,14 +15,27 @@
 import sqlalchemy as sa  # noqa
 from alembic import op  # noqa
 from jupyterhub.orm import JSONDict  # noqa
+from sqlalchemy.dialects import postgresql  # noqa
 
 
 def upgrade():
     op.create_table(
         "images",
         sa.Column("uid", sa.Unicode(36)),
-        sa.Column("name", sa.Unicode(4096)),
-        sa.Column("metadata", JSONDict, nullable=True),
+        sa.Column("name", sa.Unicode(4096), nullable=False),
+        sa.Column(
+            "status",
+            postgresql.ENUM(
+                "success",
+                "building",
+                "failed",
+                name="build_status_enum",
+                create_type=True,
+            ),
+            nullable=False,
+        ),
+        sa.Column("log", sa.UnicodeText()),
+        sa.Column("metadata", JSONDict, nullable=False),
     )
 
 

tljh_repo2docker/app.py

@@ -16,7 +16,9 @@
 from tljh_repo2docker.binderhub_builder import BinderHubBuildHandler
 
 from .builder import BuildHandler
-from .dbutil import async_session_context_factory, sync_to_async_url, upgrade_if_needed
+from .database.manager import ImagesDatabaseManager
+from .dbutil import (async_session_context_factory, sync_to_async_url,
+                     upgrade_if_needed)
 from .environments import EnvironmentsHandler
 from .logs import LogsHandler
 from .servers import ServersHandler
@@ -197,6 +199,8 @@ def init_settings(self) -> tp.Dict:
         )
         if hasattr(self, "db_context"):
             settings["db_context"] = self.db_context
+        if hasattr(self, "image_db_manager"):
+            settings["image_db_manager"] = self.image_db_manager
         return settings
 
     def init_handlers(self) -> tp.List:
@@ -277,6 +281,8 @@ def init_db(self):
             self.log.error("Failed to connect to db: %s", db_log_url)
             self.log.debug("Database error was:", exc_info=True)
 
+        self.image_db_manager = ImagesDatabaseManager()
+
     def make_app(self) -> web.Application:
         """Create the tornado web application.
         Returns:

tljh_repo2docker/binderhub_builder.py

@@ -1,10 +1,11 @@
 import json
+from urllib.parse import quote
 
 import requests
+from jupyterhub.utils import url_path_join
 from tornado import web
+
 from .base import BaseHandler, require_admin_role
-from urllib.parse import quote
-from jupyterhub.utils import url_path_join
 
 
 class BinderHubBuildHandler(BaseHandler):
@@ -28,12 +29,12 @@ async def post(self):
         url = url_path_join(binder_url, "build", provider, quoted_repo, ref)
 
         params = {"build_only": "true"}
-        
+
         async with self.client.stream("GET", url, params=params, timeout=None) as r:
             async for line in r.aiter_lines():
                 print("proviDDDDDDDDDDDDDDDDder", line)
                 # if line.startswith("data:"):
-                    # print(line.split(":", 1)[1])
+                # print(line.split(":", 1)[1])
 
         self.set_status(200)
         self.set_header("content-type", "application/json")

tljh_repo2docker/database/manager.py

@@ -0,0 +1,186 @@
+import logging
+from typing import List, Type, Union
+
+import sqlalchemy as sa
+from pydantic import UUID4
+from sqlalchemy.exc import IntegrityError, SQLAlchemyError
+from sqlalchemy.ext.asyncio import AsyncSession
+from tornado.web import HTTPError
+
+from .model import DockerImageSQL
+from .schemas import (DockerImageCreateSchema, DockerImageOutSchema,
+                      DockerImageUpdateSchema)
+
+
+class ImagesDatabaseManager:
+
+    @property
+    def _table(self) -> Type[DockerImageSQL]:
+        return DockerImageSQL
+
+    @property
+    def _schema_out(self) -> Type[DockerImageOutSchema]:
+        return DockerImageOutSchema
+
+    async def create(
+        self, db: AsyncSession, obj_in: DockerImageCreateSchema
+    ) -> DockerImageOutSchema:
+        """
+        Create one resource.
+
+        ```sql
+        INSERT INTO resource
+        VALUES (...)
+        ```
+
+        Args:
+            db: An asyncio version of SQLAlchemy session.
+            obj_in: An object containing the resource instance to create
+
+        Returns:
+            The created resource instance on success.
+
+        Raises:
+            DatabaseError: If `db.commit()` failed.
+        """
+        entry = self._table(**obj_in.model_dump())
+
+        db.add(entry)
+
+        try:
+            await db.commit()
+            # db.refresh(entry)
+        except IntegrityError as e:
+            logging.error(f"create: {e}")
+            raise HTTPError(409, "That resource already exists.")
+        except SQLAlchemyError as e:
+            logging.error(f"create: {e}")
+            raise e
+
+        return self._schema_out.model_validate(entry)
+
+    async def read(
+        self, db: AsyncSession, uid: UUID4
+    ) -> Union[DockerImageOutSchema, None]:
+        """
+        Get one resource by uid.
+
+        ```sql
+        SELECT *
+        FROM %tablename
+        WHERE uid=%s
+        ```
+
+        Args:
+            db: An asyncio version of SQLAlchemy session.
+            uid: The primary key of the resource to retrieve.
+
+        Returns:
+            The first resource instance found, `None` if no instance retrieved.
+        """
+        if entry := await db.get(self._table, uid):
+            return self._schema_out.model_validate(entry)
+        return None
+
+    async def read_many(
+        self, db: AsyncSession, uids: List[UUID4]
+    ) -> List[DockerImageOutSchema]:
+        """
+        Get multiple resources.
+
+        ```sql
+        SELECT *
+        FROM %tablename
+        WHERE uid=_in(%s...)
+        ```
+
+        Args:
+            db: An asyncio version of SQLAlchemy session.
+            uids: The primary keys of the resources to retrieve.
+
+        Returns:
+            The list of resources retrieved.
+        """
+        resources = (
+            await db.execute(sa.select(self._table).where(self._table.uid.in_(uids)))
+        ).scalars()
+        return [self._schema_out.model_validate(r) for r in resources]
+
+    async def update(
+        self, db: AsyncSession, obj_in: DockerImageUpdateSchema, optimistic: bool = True
+    ) -> Union[DockerImageOutSchema, None]:
+        """
+        Update one object.
+
+        ```sql
+        UPDATE %tablename
+        SET (...), updated_at=now()
+        WHERE uid=%s
+        ```
+
+        Args:
+            db: An asyncio version of SQLAlchemy session.
+            obj_in: A model containing values to update
+            optimistic: If `True`, assert the new model instance to be
+            `**{**obj_db.dict(), **obj_in.dict()}`
+
+        Returns:
+            The updated model instance on success, `None` if it does not exist
+            yet in database.
+
+        Raises:
+            DatabaseError: If `db.commit()` failed.
+        """
+        if not (obj_db := await self.read(db=db, uid=obj_in.uid)):
+            await self.create(db, obj_in)
+
+        update_data = obj_in.model_dump()
+
+        await db.execute(
+            sa.update(self._table)
+            .where(self._table.uid == obj_in.uid)
+            .values(**update_data)
+        )
+
+        try:
+            await db.commit()
+        except SQLAlchemyError as e:
+            logging.error(f"update: {e}")
+            raise e
+
+        if optimistic:
+            for field in update_data:
+                setattr(obj_db, field, update_data[field])
+            return self._schema_out.model_validate(obj_db)
+
+        return await self.read(db=db, uid=obj_in.uid)
+
+    async def delete(self, db: AsyncSession, uid: UUID4) -> bool:
+        """
+        Delete one object.
+
+        ```sql
+        DELETE
+        FROM %tablename
+        WHERE uid=%s
+        ```
+
+        Args:
+            db: An asyncio version of SQLAlchemy session.
+            uid: The primary key of the resource to delete.
+
+        Returns:
+            bool: `True` if the object has been deleted, `False` otherwise.
+
+        Raises:
+            DatabaseError: If `db.commit()` failed.
+        """
+        results = await db.execute(sa.delete(self._table).where(self._table.uid == uid))
+
+        try:
+            await db.commit()
+        except SQLAlchemyError as e:
+            logging.error(f"delete: {e}")
+            raise e
+
+        return results.rowcount == 1

tljh_repo2docker/database/model.py

@@ -0,0 +1,42 @@
+import uuid
+
+from jupyterhub.orm import JSONDict
+from sqlalchemy import Column, String, Text
+from sqlalchemy.dialects.postgresql import ENUM, UUID
+from sqlalchemy.ext.declarative import DeclarativeMeta, declarative_base
+
+from .schemas import BuildStatusType
+
+BaseSQL: DeclarativeMeta = declarative_base()
+
+
+class DockerImageSQL(BaseSQL):
+    """
+    SQLAlchemy image table definition.
+    """
+
+    __tablename__ = "images"
+
+    uid = Column(
+        UUID(as_uuid=True),
+        primary_key=True,
+        default=uuid.uuid4,
+    )
+
+    name = Column(String(length=4096), unique=False, nullable=False)
+
+    status = Column(
+        ENUM(
+            BuildStatusType,
+            name="build_status_enum",
+            create_type=False,
+            values_callable=lambda enum: [e.value for e in enum],
+        ),
+        nullable=False,
+    )
+
+    log = Column(Text)
+
+    metadata = Column(JSONDict, default={})
+
+    __mapper_args__ = {"eager_defaults": True}

tljh_repo2docker/database/schemas.py

@@ -0,0 +1,40 @@
+from enum import Enum
+
+from pydantic import UUID4, BaseModel
+
+
+class BuildStatusType(str, Enum):
+    SUCCESS = "success"
+    BUILDING = "building"
+    FAILED = "failed"
+
+
+class ImageMetadataType(BaseModel):
+    label: str
+    repo: str
+    ref: str
+    cpu: str
+    memory: str
+
+
+class DockerImageCreateSchema(BaseModel):
+    uid: UUID4
+    name: str
+    status: BuildStatusType
+    log: str
+    metadata: ImageMetadataType
+
+    class Config:
+        use_enum_values = True
+
+
+class DockerImageUpdateSchema(DockerImageCreateSchema):
+    pass
+
+
+class DockerImageOutSchema(DockerImageCreateSchema):
+
+    class Config:
+        use_enum_values = True
+        from_attributes = True
+        orm_mode = True

tljh_repo2docker/dbutil.py

@@ -13,7 +13,8 @@
 import alembic.config
 from alembic.script import ScriptDirectory
 from sqlalchemy import create_engine, inspect, text
-from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
+from sqlalchemy.ext.asyncio import (AsyncSession, async_sessionmaker,
+                                    create_async_engine)
 
 HERE = Path(__file__).parent.resolve()
 ALEMBIC_DIR = HERE / "alembic"

tljh_repo2docker/environments.py

@@ -28,7 +28,7 @@ async def get(self):
             default_cpu_limit=self.settings.get("default_cpu_limit"),
             machine_profiles=self.settings.get("machine_profiles", []),
             repo_providers=self.settings.get("repo_providers", None),
-            use_binderhub=use_binderhub
+            use_binderhub=use_binderhub,
         )
         if isawaitable(result):
             self.write(await result)