Upgrade pydantic to 2.x

Author: erikamov

airflow/dags/airtable_loader_v2/generate_gtfs_download_configs.py

@@ -19,7 +19,7 @@
     GTFSDownloadConfigExtract,
     get_fs,
 )
-from pydantic import ValidationError
+from pydantic.v1 import ValidationError
 
 
 def gtfs_datasets_to_extract_configs(

airflow/dags/download_gtfs_schedule_v2/download_schedule_feeds.py

@@ -24,7 +24,7 @@
     get_fs,
     get_latest,
 )
-from pydantic import validator
+from pydantic.v1 import validator
 from requests.exceptions import HTTPError
 
 GTFS_FEED_LIST_ERROR_THRESHOLD = 0.95

airflow/dags/sync_ntd_data_xlsx/scrape_ntd_xlsx_urls.py

@@ -7,7 +7,7 @@
 
 import requests
 from bs4 import BeautifulSoup
-from pydantic import HttpUrl, ValidationError, parse_obj_as
+from pydantic.v1 import HttpUrl, ValidationError, parse_obj_as
 from requests.adapters import HTTPAdapter
 from urllib3.util.retry import Retry
 

airflow/plugins/calitp_data_infra/__init__.py

@@ -0,0 +1,74 @@
+import os
+from typing import Mapping
+
+from google.cloud import secretmanager  # type: ignore
+
+
+def get_gcp_project_id() -> str:
+    """
+    Get the GCP project ID from the environment variable. Try the following
+    environment variables in order:
+    1. GOOGLE_CLOUD_PROJECT
+    2. GCP_PROJECT
+    3. PROJECT_ID
+
+    Cloud Composer should set at least one of these environment variables. See
+    https://cloud.google.com/composer/docs/composer-3/set-environment-variables
+    """
+    try:
+        return os.environ["GOOGLE_CLOUD_PROJECT"]
+    except KeyError:
+        pass
+
+    try:
+        return os.environ["GCP_PROJECT"]
+    except KeyError:
+        pass
+
+    try:
+        return os.environ["PROJECT_ID"]
+    except KeyError:
+        raise ValueError(
+            "GCP project not set in environment variables. "
+            "Please set GCP_PROJECT, GOOGLE_CLOUD_PROJECT, "
+            "or PROJECT_ID."
+        )
+
+
+def get_secret_by_name(
+    name: str,
+    project: str = "",
+    client=secretmanager.SecretManagerServiceClient(),
+) -> str:
+    project = project or get_gcp_project_id()
+
+    version = f"projects/{project}/secrets/{name}/versions/latest"
+    response = client.access_secret_version(name=version)
+    return response.payload.data.decode("UTF-8").strip()
+
+
+def get_secrets_by_label(
+    label: str,
+    project: str = "",
+    client=secretmanager.SecretManagerServiceClient(),
+) -> Mapping[str, str]:
+    secret_values = {}
+    project = project or get_gcp_project_id()
+
+    # once we get on at least 2.0.0 of secret manager, we can filter server-side
+    for secret in client.list_secrets(parent=f"projects/{project}"):
+        if label in secret.labels:
+            version = f"{secret.name}/versions/latest"
+            response = client.access_secret_version(name=version)
+            secret_values[secret.name.split("/")[-1]] = response.payload.data.decode(
+                "UTF-8"
+            ).strip()
+
+    return secret_values
+
+
+if __name__ == "__main__":
+    print("loading secrets...")
+    get_secret_by_name("BEAR_TRANSIT_KEY")
+    print(get_secrets_by_label("gtfs_schedule").keys())
+    print(get_secrets_by_label("gtfs_rt").keys())