Upgrade to pydantic 2

pctasks/cli/requirements.txt

@@ -6,6 +6,9 @@ aiohttp==3.9.5
 aiosignal==1.3.1
     # via aiohttp
     # from https://pypi.org/simple
+annotated-types==0.7.0
+    # via pydantic
+    # from https://pypi.org/simple
 async-timeout==4.0.3
     # via aiohttp
     # from https://pypi.org/simple
@@ -163,10 +166,17 @@ pyasn1-modules==0.4.0
 pycparser==2.22
     # via cffi
     # from https://pypi.org/simple
-pydantic==1.10.15
+pydantic==2.10.6
     # via
     #   pctasks-core
     #   planetary-computer
+    #   pydantic-settings
+    # from https://pypi.org/simple
+pydantic-core==2.27.2
+    # via pydantic
+    # from https://pypi.org/simple
+pydantic-settings==2.8.1
+    # via pctasks-core
     # from https://pypi.org/simple
 pyjwt==2.8.0
     # via msal
@@ -191,7 +201,9 @@ python-dateutil==2.8.2
     #   strictyaml
     # from https://pypi.org/simple
 python-dotenv==1.0.1
-    # via planetary-computer
+    # via
+    #   planetary-computer
+    #   pydantic-settings
     # from https://pypi.org/simple
 pytz==2024.1
     # via planetary-computer
@@ -225,14 +237,15 @@ stac-validator==3.3.2
 strictyaml==1.7.3
     # via pctasks-core
     # from https://pypi.org/simple
-typing-extensions==4.11.0
+typing-extensions==4.12.2
     # via
     #   azure-core
     #   azure-cosmos
     #   azure-data-tables
     #   azure-storage-blob
     #   azure-storage-queue
     #   pydantic
+    #   pydantic-core
     # from https://pypi.org/simple
 urllib3==2.2.1
     # via requests

pctasks/client/pctasks/client/client.py

@@ -174,7 +174,7 @@ def _yield_page_results(
                 route,
                 params=params,
             )
-            result = record_list_response_type.parse_obj(resp)
+            result = record_list_response_type.model_validate(resp)
             yield from result.records
             page_count += 1
             page_token = result.next_page_token
@@ -285,7 +285,7 @@ def get_workflow(self, workflow_id: str) -> Optional[WorkflowRecord]:
             result = self._call_api(
                 "GET", WORKFLOW_ROUTE.format(workflow_id=workflow_id)
             )
-            return WorkflowRecordResponse.parse_obj(result).record
+            return WorkflowRecordResponse.model_validate(result).record
         except HTTPError as e:
             if e.response.status_code == 404:
                 return None
@@ -369,7 +369,7 @@ def submit_workflow(
 
         workflow_def = workflow.definition
 
-        request = request.copy()
+        request = request.model_copy()
 
         # Ensure arguments
         request.args = self.settings.add_default_args(
@@ -452,7 +452,7 @@ def get_workflow_run(
             route += f"?dataset={dataset_id}"
         try:
             result = self._call_api("GET", route)
-            return WorkflowRunRecordResponse.parse_obj(result).record
+            return WorkflowRunRecordResponse.model_validate(result).record
         except HTTPError as e:
             if e.response.status_code == 404:
                 return None
@@ -497,7 +497,7 @@ def get_job_partition_run(
         )
         try:
             result = self._call_api("GET", route)
-            return JobPartitionRunRecordResponse.parse_obj(result).record
+            return JobPartitionRunRecordResponse.model_validate(result).record
         except HTTPError as e:
             if e.response.status_code == 404:
                 return None

pctasks/client/pctasks/client/profile/commands.py

@@ -141,7 +141,7 @@ def set_profile(ctx: click.Context, profile: str) -> None:
         rprint(f"[red]Profile [bold]{profile}[/bold] does not exists[/red]")
         ctx.exit(1)
 
-    profile_only_config = settings_config.copy(update={"settings_file": None})
+    profile_only_config = settings_config.model_copy(update={"settings_file": None})
     profile_settings_file = profile_only_config.get_settings_file()
     if not Path(profile_settings_file).exists():
         raise click.UsageError(
@@ -160,7 +160,7 @@ def show_profile(ctx: click.Context, profile: str) -> None:
         rprint(f"[red]Profile [bold]{profile}[/bold] does not exists[/red]")
         ctx.exit(1)
 
-    profile_only_config = settings_config.copy(update={"settings_file": None})
+    profile_only_config = settings_config.model_copy(update={"settings_file": None})
     yaml_txt = profile_only_config.get_settings_file().read_text()
     console = Console()
     console.print(Syntax(yaml_txt, "yaml"))

pctasks/client/pctasks/client/settings.py

@@ -1,7 +1,7 @@
 from typing import Dict, Optional
 from urllib.parse import urlparse
 
-from pydantic import validator
+from pydantic import field_validator
 
 from pctasks.core.models.workflow import WorkflowDefinition
 from pctasks.core.settings import PCTasksSettings
@@ -21,7 +21,7 @@ def section_name(cls) -> str:
     default_args: Optional[Dict[str, str]] = None
     default_page_size: int = DEFAULT_PAGE_SIZE
 
-    @validator("endpoint")
+    @field_validator("endpoint", mode="after")
     def _validate_endpoint(cls, v: str) -> str:
         try:
             parsed = urlparse(v)

pctasks/client/pctasks/client/workflow/commands.py

@@ -264,7 +264,7 @@ def cli_handle_workflow(
         workflow_id = workflow_def.workflow_id
 
     else:
-        workflow_def = workflow_def.copy(update={"workflow_id": workflow_id})
+        workflow_def = workflow_def.model_copy(update={"workflow_id": workflow_id})
 
     if not client:
         client = PCTasksClient(settings=ClientSettings.from_context(ctx.obj))

pctasks/client/pctasks/client/workflow/template.py

@@ -24,7 +24,7 @@ def template_workflow_dict(
     base_path = base_path or Path(".")
     workflow_dict = LocalTemplater(base_path).template_dict(workflow_dict)
 
-    return WorkflowDefinition.parse_obj(workflow_dict)
+    return WorkflowDefinition.model_validate(workflow_dict)
 
 
 def template_workflow_contents(

pctasks/client/pyproject.toml

@@ -26,7 +26,7 @@ classifiers = [
 dependencies = [
     "pctasks.core @ {root:parent:uri}/core",
     "pctasks.cli @ {root:parent:uri}/cli",
-    "pydantic[dotenv]>=1.8,<2.0.0",
+    "pydantic[dotenv]>=2.0.0",
     "rich>=11.2.0",
 ]
 

pctasks/client/requirements.txt

@@ -6,6 +6,9 @@ aiohttp==3.9.5
 aiosignal==1.3.1
     # via aiohttp
     # from https://pypi.org/simple
+annotated-types==0.7.0
+    # via pydantic
+    # from https://pypi.org/simple
 async-timeout==4.0.3
     # via aiohttp
     # from https://pypi.org/simple
@@ -173,11 +176,18 @@ pyasn1-modules==0.4.0
 pycparser==2.22
     # via cffi
     # from https://pypi.org/simple
-pydantic==1.10.15
+pydantic==2.10.6
     # via
     #   pctasks-client (./pctasks/client/pyproject.toml)
     #   pctasks-core
     #   planetary-computer
+    #   pydantic-settings
+    # from https://pypi.org/simple
+pydantic-core==2.27.2
+    # via pydantic
+    # from https://pypi.org/simple
+pydantic-settings==2.8.1
+    # via pctasks-core
     # from https://pypi.org/simple
 pygments==2.17.2
     # via rich
@@ -207,7 +217,7 @@ python-dateutil==2.8.2
 python-dotenv==1.0.1
     # via
     #   planetary-computer
-    #   pydantic
+    #   pydantic-settings
     # from https://pypi.org/simple
 pytz==2024.1
     # via planetary-computer
@@ -244,14 +254,15 @@ stac-validator==3.3.2
 strictyaml==1.7.3
     # via pctasks-core
     # from https://pypi.org/simple
-typing-extensions==4.11.0
+typing-extensions==4.12.2
     # via
     #   azure-core
     #   azure-cosmos
     #   azure-data-tables
     #   azure-storage-blob
     #   azure-storage-queue
     #   pydantic
+    #   pydantic-core
     # from https://pypi.org/simple
 urllib3==2.2.1
     # via requests

pctasks/client/tests/test_client.py

@@ -9,6 +9,8 @@
 from pctasks.dev.blob import get_azurite_code_storage
 from pctasks.dev.test_utils import assert_workflow_is_successful
 
+# from pctasks.dev.test_utils import assert_workflow_is_successful
+
 HERE = pathlib.Path(__file__).parent
 
 

pctasks/client/tests/test_template.py

@@ -23,7 +23,7 @@ def test_local_file_template():
     yaml_dict = yaml.safe_load(yaml_str)
     templated_dict = LocalTemplater().template_dict(yaml_dict)
 
-    data = IngestCollectionsInput.parse_obj(templated_dict)
+    data = IngestCollectionsInput.model_validate(templated_dict)
 
     assert data.collections
     assert data.collections[0]["id"] == "test-collection"

pctasks/core/pctasks/core/activity.py

@@ -34,13 +34,13 @@ def _func(msg: str) -> str:
         if "msg" not in msg_dict:
             raise ValueError(f"Missing 'msg' in message: {msg_dict}")
 
-        msg_model = model_class.parse_obj(msg_dict["msg"])
+        msg_model = model_class.model_validate(msg_dict["msg"])
         event_tag_msg: Optional[str] = None
         if event_tag:
             event_tag_msg = event_tag(msg_model)
 
         activity_msg: ActivityMessage[T] = ActivityMessage(
-            run_record_id=RunRecordId.parse_obj(msg_dict.get("run_record_id")),
+            run_record_id=RunRecordId.model_validate(msg_dict.get("run_record_id")),
             msg=msg_model,
         )
 

pctasks/core/pctasks/core/cosmos/container.py

@@ -128,7 +128,7 @@ def item_from_model(self, model: T) -> Dict[str, Any]:
 
     def model_from_item(self, model_type: Type[T], item: Dict[str, Any]) -> T:
         """Transform a cosmosdb item (dict) into a model."""
-        return model_type.parse_obj(item)
+        return model_type.model_validate(item)
 
     def _prepare_put_item(self, model: T) -> Dict[str, Any]:
         # Set created or updated time

pctasks/core/pctasks/core/cosmos/settings.py

@@ -1,13 +1,14 @@
 import os
 import re
-from typing import Any, Dict, Optional
+from typing import Optional
 from urllib.parse import urlparse
 
 import azure.identity.aio
 from azure.cosmos import CosmosClient
 from azure.cosmos.aio import CosmosClient as AsyncCosmosClient
 from azure.identity import DefaultAzureCredential
-from pydantic import validator
+from pydantic import field_validator, model_validator
+from typing_extensions import Self
 
 from pctasks.core.constants import COSMOSDB_EMULATOR_HOST_ENV_VAR
 from pctasks.core.settings import PCTasksSettings
@@ -67,11 +68,11 @@ def get_records_container_name(self) -> str:
     def get_process_item_errors_container_name(self) -> str:
         return f"{self.process_item_errors_container_name}{self.test_container_suffix}"
 
-    @validator("test_container_suffix", always=True)
+    @field_validator("test_container_suffix")
     def _validate_test_container_suffix(cls, v: Optional[str]) -> str:
         return v or ""
 
-    @validator("connection_string")
+    @field_validator("connection_string")
     def _validate_connection_string(cls, v: Optional[str]) -> Optional[str]:
         if v:
             if not re.search(r"AccountEndpoint=(.*?);", v):
@@ -80,15 +81,14 @@ def _validate_connection_string(cls, v: Optional[str]) -> Optional[str]:
                 raise ValueError("Cannot find AccountKey in connection_string")
         return v
 
-    @validator("key", always=True)  # Validates all connection properties (last defined)
-    def _validate_key(cls, v: Optional[str], values: Dict[str, Any]) -> Optional[str]:
-        if v:
-            if values.get("connection_string"):
-                raise ValueError("Cannot set both key and connection_string")
-            if not values.get("url"):
-                raise ValueError("Must set url when setting key")
+    @model_validator(mode="after")  # Validates all connection properties (last defined)
+    def _key_or_connection_string(self) -> Self:
+        if self.key and self.connection_string:
+            raise ValueError("Cannot set both key and connection_string")
+        if self.key and not self.url:
+            raise ValueError("Must set url when setting key")
 
-        return v
+        return self
 
     def get_cosmosdb_url(self) -> str:
         if self.connection_string:

pctasks/core/pctasks/core/models/activity.py

@@ -7,6 +7,6 @@
 T = TypeVar("T", bound=BaseModel)
 
 
-class ActivityMessage(Generic[T], PCBaseModel):
+class ActivityMessage(PCBaseModel, Generic[T]):
     run_record_id: RunRecordId
     msg: T

pctasks/core/pctasks/core/models/base.py

@@ -14,12 +14,12 @@ class PCBaseModel(BaseModel):
     def dict(self, **kwargs: Any) -> Dict[str, Any]:
         kwargs.setdefault("by_alias", True)
         kwargs.setdefault("exclude_none", True)
-        return super().dict(**kwargs)
+        return super().model_dump(**kwargs)
 
     def json(self, **kwargs: Any) -> str:
         kwargs.setdefault("by_alias", True)
         kwargs.setdefault("exclude_none", True)
-        return super().json(**kwargs)
+        return super().model_dump_json(**kwargs)
 
     def to_json(self, *args: Any, **kwargs: Any) -> str:
         """Passed through to .json()
@@ -58,9 +58,9 @@ def _prop_sort(field_order: List[str], d: Dict[str, Any]) -> Dict[str, Any]:
     def from_yaml(cls: Type[T], yaml_str: str, section: Optional[str] = None) -> T:
         return model_from_yaml(cls, yaml_str, section=section)
 
-    class Config:
-        exclude_none = True
-        allow_population_by_field_name = True
+    model_config = {
+        "populate_by_name": True,
+    }
 
 
 class RunRecordId(PCBaseModel):

pctasks/core/pctasks/core/models/config.py

@@ -1,6 +1,6 @@
 from typing import Dict, List, Optional
 
-from pydantic import validator
+from pydantic import field_validator
 
 from pctasks.core.models.base import PCBaseModel
 from pctasks.core.storage.blob import BlobStorage, BlobUri
@@ -59,15 +59,15 @@ class ImageConfig(PCBaseModel):
     environment: Optional[List[str]] = None
     tags: Optional[List[str]] = None
 
-    @validator("environment")
+    @field_validator("environment")
     def _environment_validator(cls, v: Optional[List[str]]) -> Optional[List[str]]:
         if v:
             for env in v:
                 if not len(env.split("=")) == 2:
                     raise ValueError(f"Environment entry {env} is invalid.")
         return v
 
-    @validator("tags")
+    @field_validator("tags")
     def _tags_validator(cls, v: Optional[List[str]]) -> Optional[List[str]]:
         if v:
             for tag in v: