Merge pull request #4619 from fedspendingtransparency/ftr/dev-14236-pydantic-upgrade

[DEV-14236] Pydantic upgrade

Author: aguest-kc

pyproject.toml

@@ -52,7 +52,8 @@ dependencies = [
     "psutil==5.9.*",
     "psycopg>=3.3.3",
     "py-gfm==2.0.0",
-    "pydantic[dotenv]==1.9.*",
+    "pydantic==2.12",
+    "pydantic-settings>=2.13.1",
     "python-json-logger==2.0.7",
     "requests==2.31.*",
     "retrying==1.3.4",

usaspending_api/config/envs/default.py

@@ -13,13 +13,9 @@
 import pathlib
 from typing import Any, ClassVar, Union
 
-from pydantic import (
-    AnyHttpUrl,
-    BaseSettings,
-    PostgresDsn,
-    SecretStr,
-    root_validator,
-)
+from pydantic import SecretStr, model_validator
+from pydantic.networks import AnyHttpUrl, PostgresDsn
+from pydantic_settings import BaseSettings, SettingsConfigDict
 
 from usaspending_api.config.utils import (
     ENV_SPECIFIC_OVERRIDE,
@@ -134,20 +130,22 @@ def _validate_database_conf(
             )
 
             if enough_parts:
-                pg_dsn = PostgresDsn(
-                    url=None,
+                try:
+                    _port = int(values[f"{resource_conf_prefix}_PORT"])
+                except (ValueError, TypeError):
+                    _port = None
+
+                pg_dsn = PostgresDsn.build(
                     scheme=values[f"{resource_conf_prefix}_SCHEME"],
-                    user=values[f"{resource_conf_prefix}_USER"],
-                    password=values[
-                        f"{resource_conf_prefix}_PASSWORD"
-                    ].get_secret_value(),
-                    host=values[f"{resource_conf_prefix}_HOST"],
-                    port=values[f"{resource_conf_prefix}_PORT"],
-                    path=(
-                        "/" + values[f"{resource_conf_prefix}_NAME"]
-                        if values[f"{resource_conf_prefix}_NAME"]
-                        else None
+                    username=values[f"{resource_conf_prefix}_USER"],
+                    password=(
+                        values[f"{resource_conf_prefix}_PASSWORD"].get_secret_value()
+                        if isinstance(values[f"{resource_conf_prefix}_PASSWORD"], SecretStr)
+                        else values[f"{resource_conf_prefix}_PASSWORD"]
                     ),
+                    host=values[f"{resource_conf_prefix}_HOST"],
+                    port=_port,
+                    path=values.get(f"{resource_conf_prefix}_NAME"),
                 )
                 values = eval_default_factory_from_root_validator(
                     cls, values, url_conf_name, lambda: str(pg_dsn)
@@ -159,7 +157,7 @@ def _validate_database_conf(
 
     # noinspection PyMethodParameters
     # Pydantic returns a classmethod for its validators, so the cls param is correct
-    @root_validator
+    @model_validator(mode="before")
     def _DATABASE_URL_and_parts_factory(cls, values: dict[str, Any]) -> dict[str, Any]:
         """A root validator to backfill DATABASE_URL and USASPENDING_DB_* part config vars and validate that they are
         all consistent.
@@ -169,6 +167,8 @@ def _DATABASE_URL_and_parts_factory(cls, values: dict[str, Any]) -> dict[str, An
         - ALSO validates that the parts and whole string are consistent. A ``ValueError`` is thrown if found to
         be inconsistent, which will in turn raise a ``pydantic.ValidationError`` at configuration time.
         """
+        default_fields = {name: field.default for name, field in cls.model_fields.items()}
+        values = {**default_fields, **values}
         # noinspection PyArgumentList
         cls._validate_database_conf(
             cls=cls,
@@ -181,7 +181,7 @@ def _DATABASE_URL_and_parts_factory(cls, values: dict[str, Any]) -> dict[str, An
 
     # noinspection PyMethodParameters
     # Pydantic returns a classmethod for its validators, so the cls param is correct
-    @root_validator
+    @model_validator(mode="before")
     def _BROKER_DB_and_parts_factory(cls, values: dict[str, Any]) -> dict[str, Any]:
         """A root validator to backfill BROKER_DB and BROKER_DB_* part config vars and validate
         that they are all consistent.
@@ -191,6 +191,8 @@ def _BROKER_DB_and_parts_factory(cls, values: dict[str, Any]) -> dict[str, Any]:
         - ALSO validates that the parts and whole string are consistent. A ``ValueError`` is thrown if found to
         be inconsistent, which will in turn raise a ``pydantic.ValidationError`` at configuration time.
         """
+        default_fields = {name: field.default for name, field in cls.model_fields.items()}
+        values = {**default_fields, **values}
         # noinspection PyArgumentList
         cls._validate_database_conf(
             cls=cls,
@@ -203,17 +205,17 @@ def _BROKER_DB_and_parts_factory(cls, values: dict[str, Any]) -> dict[str, Any]:
 
     # ==== [Elasticsearch] ====
     # Where to connect to elasticsearch.
-    ES_HOSTNAME: str = None  # FACTORY_PROVIDED_VALUE. See below validator-factory
+    ES_HOSTNAME: str | None = None  # FACTORY_PROVIDED_VALUE. See below validator-factory
     ES_SCHEME: str = "https"
     ES_HOST: str = ENV_SPECIFIC_OVERRIDE
-    ES_PORT: str = None
-    ES_USER: str = None
-    ES_PASSWORD: SecretStr = None
-    ES_NAME: str = None
+    ES_PORT: str | None = None
+    ES_USER: str | None = None
+    ES_PASSWORD: SecretStr | None = None
+    ES_NAME: str | None = None
 
     # noinspection PyMethodParameters
     # Pydantic returns a classmethod for its validators, so the cls param is correct
-    @root_validator
+    @model_validator(mode="before")
     def _ES_HOSTNAME_and_parts_factory(cls, values: dict[str, Any]) -> dict[str, Any]:
         """A root validator to backfill ES_HOSTNAME and ES_* part config vars and validate that they are
         all consistent.
@@ -223,6 +225,8 @@ def _ES_HOSTNAME_and_parts_factory(cls, values: dict[str, Any]) -> dict[str, Any
         - ALSO validates that the parts and whole string are consistent. A ``ValueError`` is thrown if found to
         be inconsistent, which will in turn raise a ``pydantic.ValidationError`` at configuration time.
         """
+        default_fields = {name: field.default for name, field in cls.model_fields.items()}
+        values = {**default_fields, **values}
         # noinspection PyArgumentList
         cls._validate_http_url(
             cls=cls,
@@ -251,9 +255,7 @@ def _validate_http_url(
         # - it should take precedence
         # - its values will be used to backfill any missing URL parts stored as separate config vars
         if is_full_url_provided:
-            values = backfill_url_parts_config(
-                cls, url_conf_name, resource_conf_prefix, values
-            )
+            values = backfill_url_parts_config(cls, url_conf_name, resource_conf_prefix, values)
 
         # If the full URL config is not provided, try to build-it-up from provided parts, then set the full URL
         if not is_full_url_provided:
@@ -268,21 +270,16 @@ def _validate_http_url(
 
             if enough_parts:
                 http_url = AnyHttpUrl(
-                    url=None,
                     scheme=values[f"{resource_conf_prefix}_SCHEME"],
-                    user=values[f"{resource_conf_prefix}_USER"],
+                    username=values[f"{resource_conf_prefix}_USER"],
                     password=(
                         values[f"{resource_conf_prefix}_PASSWORD"].get_secret_value()
                         if values[f"{resource_conf_prefix}_PASSWORD"]
                         else None
                     ),
                     host=values[f"{resource_conf_prefix}_HOST"],
                     port=values[f"{resource_conf_prefix}_PORT"],
-                    path=(
-                        "/" + values[f"{resource_conf_prefix}_NAME"]
-                        if values[f"{resource_conf_prefix}_NAME"]
-                        else None
-                    ),
+                    path=values.get(f"{resource_conf_prefix}_NAME"),
                 )
                 values = eval_default_factory_from_root_validator(
                     cls, values, url_conf_name, lambda: str(http_url)
@@ -298,7 +295,7 @@ def _validate_http_url(
     # Those clusters are the only place we currently need this variable,
     # If you write code that depends on this config, make sure you
     # set BRANCH as an environment variable on your machine
-    BRANCH: str = os.environ.get("BRANCH")
+    BRANCH: str | None = os.environ.get("BRANCH")
 
     # SPARK_SCHEDULER_MODE = "FAIR"  # if used with weighted pools, could allow round-robin tasking of simultaneous jobs
     # TODO: have to deal with this if really wanting balanced (FAIR) task execution
@@ -361,10 +358,10 @@ def _validate_http_url(
     AWS_ACCESS_KEY: SecretStr = ENV_SPECIFIC_OVERRIDE
     AWS_SECRET_KEY: SecretStr = ENV_SPECIFIC_OVERRIDE
     # Setting AWS_PROFILE to None so boto3 doesn't try to pick up the placeholder string as an actual profile to find
-    AWS_PROFILE: str = None  # USER_SPECIFIC_OVERRIDE
-    SPARK_S3_BUCKET: str = os.environ.get("SPARK_S3_BUCKET")
-    BULK_DOWNLOAD_S3_BUCKET_NAME: str = os.environ.get("BULK_DOWNLOAD_S3_BUCKET_NAME")
-    DATABASE_DOWNLOAD_S3_BUCKET_NAME: str = os.environ.get(
+    AWS_PROFILE: str | None = None  # USER_SPECIFIC_OVERRIDE
+    SPARK_S3_BUCKET: str | None = os.environ.get("SPARK_S3_BUCKET")
+    BULK_DOWNLOAD_S3_BUCKET_NAME: str | None = os.environ.get("BULK_DOWNLOAD_S3_BUCKET_NAME")
+    DATABASE_DOWNLOAD_S3_BUCKET_NAME: str | None = os.environ.get(
         "DATABASE_DOWNLOAD_S3_BUCKET_NAME"
     )
     DELTA_LAKE_S3_PATH: str = "data/delta"  # path within SPARK_S3_BUCKET where Delta output data will accumulate
@@ -380,9 +377,8 @@ def _validate_http_url(
     COVID19_DOWNLOAD_README_OBJECT_KEY: str = (
         f"files/{COVID19_DOWNLOAD_README_FILE_NAME}"
     )
-
-    class Config:
-        pass
-        # supporting use of a user-provided (ang git-ignored) .env file for overrides
-        env_file = str(_PROJECT_ROOT_DIR / ".env")
-        env_file_encoding = "utf-8"
+    model_config = SettingsConfigDict(
+        env_file=str(_PROJECT_ROOT_DIR / ".env"),
+        env_file_encoding="utf-8",
+        extra="allow",
+    )

usaspending_api/config/envs/local.py

@@ -7,14 +7,15 @@
 # - Set config variables to DefaultConfig.USER_SPECIFIC_OVERRIDE where there is expected to be a
 #   user-provided a config value for a variable (e.g. in the ../.env file)
 ########################################################################################################################
-from typing import ClassVar
+from typing import Any, ClassVar
 
-from pydantic import root_validator
+from pydantic import model_validator
 from pydantic.types import SecretStr
-from usaspending_api.config.envs.default import DefaultConfig, _PROJECT_ROOT_DIR
+
+from usaspending_api.config.envs.default import _PROJECT_ROOT_DIR, DefaultConfig
 from usaspending_api.config.utils import (
-    USER_SPECIFIC_OVERRIDE,
     FACTORY_PROVIDED_VALUE,
+    USER_SPECIFIC_OVERRIDE,
     eval_default_factory_from_root_validator,
 )
 
@@ -90,20 +91,24 @@ class LocalConfig(DefaultConfig):
     USE_AWS: bool = False
     AWS_ACCESS_KEY: SecretStr = MINIO_ACCESS_KEY
     AWS_SECRET_KEY: SecretStr = MINIO_SECRET_KEY
-    AWS_PROFILE: str = None
+    AWS_PROFILE: str | None = None
     AWS_REGION: str = ""
     SPARK_S3_BUCKET: str = "data"
     BULK_DOWNLOAD_S3_BUCKET_NAME: str = "bulk-download"
-    DATABASE_DOWNLOAD_S3_BUCKET_NAME = "dti-usaspending-db"
+    DATABASE_DOWNLOAD_S3_BUCKET_NAME: str = "dti-usaspending-db"
 
     # Since this config values is built by composing others, we want to late/lazily-evaluate their values,
     # in case the declared value is overridden by a shell env var or .env file value
-    AWS_S3_ENDPOINT: str = FACTORY_PROVIDED_VALUE  # See below validator-based factory
+    AWS_S3_ENDPOINT: str | None = FACTORY_PROVIDED_VALUE  # See below validator-based factory
+
+    @model_validator(mode="before")
+    def _AWS_S3_ENDPOINT_factory(cls, values: dict[str, Any]) -> dict[str, Any]:
+        # Merge defaults into values
+        default_fields = {name: field.default for name, field in cls.model_fields.items()}
+        merged_values = {**default_fields, **values}
 
-    @root_validator
-    def _AWS_S3_ENDPOINT_factory(cls, values):
-        def factory_func():
-            return values["MINIO_HOST"] + ":" + values["MINIO_PORT"]
+        def factory_func() -> str:
+            return merged_values["MINIO_HOST"] + ":" + merged_values["MINIO_PORT"]
 
         return eval_default_factory_from_root_validator(cls, values, "AWS_S3_ENDPOINT", factory_func)