Migrate datajunction-server to pydantic v2

Author: shangyian

datajunction-server/datajunction_server/api/access/authentication/whoami.py

@@ -31,7 +31,7 @@ async def whoami(
     Returns the current authenticated user
     """
     user = await User.get_by_username(session, current_user.username)
-    return UserOutput.from_orm(user)
+    return UserOutput.model_validate(user, from_attributes=True)
 
 
 @router.get("/token/")

datajunction-server/datajunction_server/api/attributes.py

@@ -35,7 +35,10 @@ async def list_attributes(
     List all available attribute types.
     """
     attributes = await AttributeType.get_all(session)
-    return [AttributeTypeBase.from_orm(attr) for attr in attributes]
+    return [
+        AttributeTypeBase.model_validate(attr, from_attributes=True)
+        for attr in attributes
+    ]
 
 
 @router.post(
@@ -62,7 +65,7 @@ async def add_attribute_type(
             message=f"Attribute type `{data.name}` already exists!",
         )
     attribute_type = await AttributeType.create(session, data)
-    return AttributeTypeBase.from_orm(attribute_type)
+    return AttributeTypeBase.model_validate(attribute_type, from_attributes=True)
 
 
 async def default_attribute_types(session: AsyncSession = Depends(get_session)):

datajunction-server/datajunction_server/api/catalogs.py

@@ -36,7 +36,7 @@ async def list_catalogs(
     """
     statement = select(Catalog).options(joinedload(Catalog.engines))
     return [
-        CatalogInfo.from_orm(catalog)
+        CatalogInfo.model_validate(catalog, from_attributes=True)
         for catalog in (await session.execute(statement)).unique().scalars()
     ]
 
@@ -111,7 +111,7 @@ async def add_catalog(
     await session.commit()
     await session.refresh(catalog, ["engines"])
 
-    return CatalogInfo.from_orm(catalog)
+    return CatalogInfo.model_validate(catalog, from_attributes=True)
 
 
 @router.post(
@@ -136,7 +136,7 @@ async def add_engines_to_catalog(
     session.add(catalog)
     await session.commit()
     await session.refresh(catalog)
-    return CatalogInfo.from_orm(catalog)
+    return CatalogInfo.model_validate(catalog, from_attributes=True)
 
 
 async def list_new_engines(

datajunction-server/datajunction_server/api/collection.py

@@ -61,7 +61,7 @@ async def create_a_collection(
     await session.commit()
     await session.refresh(collection)
 
-    return CollectionInfo.from_orm(collection)
+    return CollectionInfo.model_validate(collection, from_attributes=True)
 
 
 @router.delete("/collections/{name}", status_code=HTTPStatus.NO_CONTENT)

datajunction-server/datajunction_server/api/cubes.py

@@ -159,7 +159,7 @@ async def cube_materialization_info(
         Granularity.YEAR: "0 0 1 1 *",  # Runs at midnight on January 1st every year
     }
     upsert = UpsertCubeMaterialization(
-        job=MaterializationJobTypeEnum.DRUID_CUBE,
+        job=MaterializationJobTypeEnum.DRUID_CUBE.value.name,
         strategy=(
             MaterializationStrategy.INCREMENTAL_TIME
             if temporal_partition
@@ -186,7 +186,7 @@ async def cube_materialization_info(
         metrics=cube_config.metrics,
         strategy=upsert.strategy,
         schedule=upsert.schedule,
-        job=upsert.job.name,
+        job=upsert.job.name,  # type: ignore
         measures_materializations=cube_config.measures_materializations,
         combiners=cube_config.combiners,
     )
@@ -299,7 +299,7 @@ async def get_cube_dimension_values(
             value=row[0 : count_column[0]] if count_column else row,
             count=row[count_column[0]] if count_column else None,
         )
-        for row in result.results.__root__[0].rows
+        for row in result.results.root[0].rows
     ]
     return DimensionValues(  # pragma: no cover
         dimensions=[

datajunction-server/datajunction_server/api/data.py

@@ -141,11 +141,15 @@ async def add_availability_state(
         table=data.table,
         valid_through_ts=data.valid_through_ts,
         url=data.url,
-        min_temporal_partition=data.min_temporal_partition,
-        max_temporal_partition=data.max_temporal_partition,
+        min_temporal_partition=[
+            str(part) for part in data.min_temporal_partition or []
+        ],
+        max_temporal_partition=[
+            str(part) for part in data.max_temporal_partition or []
+        ],
         partitions=[
-            partition.dict() if not isinstance(partition, Dict) else partition
-            for partition in data.partitions  # type: ignore
+            partition.model_dump() if not isinstance(partition, Dict) else partition
+            for partition in (data.partitions or [])
         ],
         categorical_partitions=data.categorical_partitions,
         temporal_partitions=data.temporal_partitions,
@@ -159,10 +163,14 @@ async def add_availability_state(
             entity_type=EntityType.AVAILABILITY,
             node=node.name,  # type: ignore
             activity_type=ActivityType.CREATE,
-            pre=AvailabilityStateBase.from_orm(old_availability).dict()
+            pre=AvailabilityStateBase.model_validate(
+                old_availability,
+            ).model_dump()
             if old_availability
             else {},
-            post=AvailabilityStateBase.from_orm(node_revision.availability).dict(),
+            post=AvailabilityStateBase.model_validate(
+                node_revision.availability,
+            ).model_dump(),
             user=current_user.username,
         ),
         session=session,
@@ -262,8 +270,8 @@ async def get_data(
     )
 
     # Inject column info if there are results
-    if result.results.__root__:  # pragma: no cover
-        result.results.__root__[0].columns = generated_sql.columns  # type: ignore
+    if result.results.root:  # pragma: no cover
+        result.results.root[0].columns = generated_sql.columns  # type: ignore
     return result
 
 
@@ -447,8 +455,8 @@ async def get_data_for_metrics(
     )
 
     # Inject column info if there are results
-    if result.results.__root__:  # pragma: no cover
-        result.results.__root__[0].columns = translated_sql.columns or []
+    if result.results.root:  # pragma: no cover
+        result.results.root[0].columns = translated_sql.columns or []
     return result
 
 

datajunction-server/datajunction_server/api/deployments.py

@@ -77,7 +77,7 @@ async def submit(self, spec: DeploymentSpec, context: DeploymentContext) -> str:
             deployment = Deployment(
                 uuid=deployment_uuid,
                 namespace=spec.namespace,
-                spec=spec.dict(),
+                spec=spec.model_dump(),
                 status=DeploymentStatus.PENDING,
                 created_by_id=context.current_user.id,
             )
@@ -109,7 +109,7 @@ async def update_status(
             deployment = await session.get(Deployment, deployment_uuid)
             deployment.status = status
             if results is not None:
-                deployment.results = [r.dict() for r in results]
+                deployment.results = [r.model_dump() for r in results]
             await session.commit()
 
     async def _run_deployment(

datajunction-server/datajunction_server/api/djsql.py

@@ -72,8 +72,8 @@ async def get_data_for_djsql(
     )
 
     # Inject column info if there are results
-    if result.results.__root__:  # pragma: no cover
-        result.results.__root__[0].columns = translated_sql.columns or []
+    if result.results.root:  # pragma: no cover
+        result.results.root[0].columns = translated_sql.columns or []
     return result
 
 

datajunction-server/datajunction_server/api/engines.py

@@ -43,7 +43,7 @@ async def list_engines(
     List all available engines
     """
     return [
-        EngineInfo.from_orm(engine)
+        EngineInfo.model_validate(engine)
         for engine in (await session.execute(select(Engine))).scalars()
     ]
 
@@ -58,7 +58,9 @@ async def get_an_engine(
     """
     Return an engine by name and version
     """
-    return EngineInfo.from_orm(await get_engine(session, name, version))
+    return EngineInfo.model_validate(
+        await get_engine(session, name, version),
+    )
 
 
 @router.post(
@@ -95,4 +97,4 @@ async def add_engine(
     await session.commit()
     await session.refresh(engine)
 
-    return EngineInfo.from_orm(engine)
+    return EngineInfo.model_validate(engine)

datajunction-server/datajunction_server/api/helpers.py

@@ -681,8 +681,8 @@ async def query_event_stream(
                 "query end state detected (%s), sending final event to the client",
                 query_next.state,
             )
-            if query_next.results.__root__:  # pragma: no cover
-                query_next.results.__root__[0].columns = columns or []
+            if query_next.results.root:  # pragma: no cover
+                query_next.results.root[0].columns = columns or []
             yield {
                 "event": "message",
                 "id": uuid.uuid4(),
@@ -879,13 +879,13 @@ def get_node_revision_materialization(
             )
             if materialization.strategy != MaterializationStrategy.INCREMENTAL_TIME:
                 info.urls = [info.urls[0]]
-            materialization_config_output = MaterializationConfigOutput.from_orm(
+            materialization_config_output = MaterializationConfigOutput.model_validate(
                 materialization,
             )
             materializations.append(
                 MaterializationConfigInfoUnified(
-                    **materialization_config_output.dict(),
-                    **info.dict(),
+                    **materialization_config_output.model_dump(),
+                    **info.model_dump(),
                 ),
             )
     return materializations