Merge branch 'main' into sync-labels

chowda/db.py

@@ -1,6 +1,11 @@
+from psycopg2.extensions import register_adapter
+from pydantic_core import Url
 from sqlalchemy import create_engine
 
 from chowda.config import DB_URL, DEBUG
+from chowda.utils import adapt_url
+
+register_adapter(Url, adapt_url)
 
 engine = create_engine(DB_URL, echo=DEBUG)
 

chowda/models.py

@@ -8,18 +8,12 @@
 from typing import Any, Dict, List, Optional
 
 from metaflow import Run, namespace
-from pydantic import AnyHttpUrl, EmailStr, stricturl
+from pydantic.networks import AnyHttpUrl, EmailStr
 from sqlalchemy import JSON, Column, DateTime, Enum
 from sqlalchemy.dialects import postgresql
-from sqlmodel import Field, Relationship, SQLModel
+from sqlmodel import AutoString, Field, Relationship, SQLModel
 from starlette.requests import Request
 
-MediaUrl = stricturl(allowed_schemes=['video', 'audio', 'text'], tld_required=False)
-"""Media url validator. Must have prefix of video, audio, or text. No TLD required.
-Example:
-    video://*
-"""
-
 
 class AppStatus(enum.Enum):
     PENDING = 'pending'
@@ -65,7 +59,7 @@ class User(SQLModel, table=True):
 
     __tablename__ = 'users'
     id: Optional[int] = Field(primary_key=True)
-    email: EmailStr = Field(index=True)
+    email: EmailStr = Field(unique=True, index=True, sa_type=AutoString)
     first_name: str = Field(min_length=3, index=True)
     last_name: str = Field(min_length=3, index=True)
 
@@ -247,7 +241,7 @@ class ClamsApp(SQLModel, table=True):
     __tablename__ = 'clams_apps'
     id: Optional[int] = Field(primary_key=True, default=None)
     name: str
-    endpoint: AnyHttpUrl
+    endpoint: AnyHttpUrl = Field(index=True, sa_type=AutoString)
     description: str
     pipelines: List['Pipeline'] = Relationship(
         back_populates='clams_apps', link_model=ClamsAppPipelineLink

chowda/utils.py

@@ -1,10 +1,16 @@
 from typing import Any, Dict, List, Set
 
+from psycopg2.extensions import QuotedString
 from pydantic import BaseModel
 from sqlalchemy.dialects.postgresql import insert
 from starlette.requests import Request
 
 
+def adapt_url(url):
+    """Adapt a Pydantic2 Url to a psycopg2 QuotedString"""
+    return QuotedString(str(url))
+
+
 def upsert(
     model: BaseModel,
     value: BaseModel,
@@ -13,10 +19,10 @@ def upsert(
     """Returns the SQLAlchemy statement to upsert a values into a table"""
     return (
         insert(model)
-        .values(value.dict())
+        .values(value.model_dump())
         .on_conflict_do_update(
             index_elements=index_elements,
-            set_=value.dict(),
+            set_=value.model_dump(),
         )
     )
 

chowda/views.py

@@ -365,7 +365,7 @@ async def duplicate_batches(
             with Session(engine) as db:
                 for batch_id in pks:
                     existing_batch = db.get(Batch, batch_id)
-                    new_batch_params = existing_batch.dict()
+                    new_batch_params = existing_batch.model_dump()
                     new_batch_params.pop('id')
                     new_batch_params['name'] += ' [COPY]'
                     new_batch = Batch(**new_batch_params)

docs/examples/upsert.ipynb

@@ -9,7 +9,7 @@
     "from sqlmodel import Session, select\n",
     "\n",
     "from chowda.models import SonyCiAsset\n",
-    "from chowda.db import engine"
+    "from chowda.db import engine\n"
    ]
   },
   {
@@ -22,7 +22,7 @@
     "    with Session(engine) as session:\n",
     "        z = session.exec(select(SonyCiAsset).where(SonyCiAsset.id == 'test'))\n",
     "\n",
-    "        x = z.first()\n"
+    "        x = z.first()"
    ]
   },
   {
@@ -56,7 +56,7 @@
     "    format='mp4',\n",
     "    thumbnails=[],\n",
     ")\n",
-    "c.dict()\n"
+    "c.model_dump()"
    ]
   },
   {
@@ -76,7 +76,7 @@
     "            )\n",
     "        )\n",
     "        session.execute(stmt)\n",
-    "        session.commit()"
+    "        session.commit()\n"
    ]
   }
  ],
@@ -96,7 +96,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.3"
+   "version": "3.11.5"
   },
   "orig_nbformat": 4
  },