diff --git a/pyproject.toml b/pyproject.toml index 2e5d670..e3f37d3 100644 --- a/pyproject.toml +++ b/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"} diff --git a/tljh_repo2docker/alembic/versions/ac1b4e7e52f3_first_migration.py b/tljh_repo2docker/alembic/versions/ac1b4e7e52f3_first_migration.py index fe7bee3..9d7b717 100644 --- a/tljh_repo2docker/alembic/versions/ac1b4e7e52f3_first_migration.py +++ b/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), ) diff --git a/tljh_repo2docker/app.py b/tljh_repo2docker/app.py index 820ff75..f0942ae 100644 --- a/tljh_repo2docker/app.py +++ b/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: diff --git a/tljh_repo2docker/binderhub_builder.py b/tljh_repo2docker/binderhub_builder.py index d6a6a38..4e799ec 100644 --- a/tljh_repo2docker/binderhub_builder.py +++ b/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") diff --git a/tljh_repo2docker/database/manager.py b/tljh_repo2docker/database/manager.py new file mode 100644 index 0000000..076ed02 --- /dev/null +++ b/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 diff --git a/tljh_repo2docker/database/model.py b/tljh_repo2docker/database/model.py new file mode 100644 index 0000000..842a87f --- /dev/null +++ b/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} diff --git a/tljh_repo2docker/database/schemas.py b/tljh_repo2docker/database/schemas.py new file mode 100644 index 0000000..bba10c8 --- /dev/null +++ b/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 diff --git a/tljh_repo2docker/dbutil.py b/tljh_repo2docker/dbutil.py index 8bd06f4..3453678 100644 --- a/tljh_repo2docker/dbutil.py +++ b/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" diff --git a/tljh_repo2docker/environments.py b/tljh_repo2docker/environments.py index b248952..312ea8f 100644 --- a/tljh_repo2docker/environments.py +++ b/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) diff --git a/tljh_repo2docker/tests/utils.py b/tljh_repo2docker/tests/utils.py index 3465104..14d9d3e 100644 --- a/tljh_repo2docker/tests/utils.py +++ b/tljh_repo2docker/tests/utils.py @@ -2,13 +2,8 @@ import json from aiodocker import Docker, DockerError -from jupyterhub.tests.utils import ( - async_requests, - auth_header, - check_db_locks, - public_host, - public_url, -) +from jupyterhub.tests.utils import (async_requests, auth_header, + check_db_locks, public_host, public_url) from jupyterhub.utils import url_path_join as ujoin from tornado.httputil import url_concat