Add ability to provide proxy for dbt Cloud connection (apache#42737)

* Add ability to provide proxy for dbt Cloud connection

* Running pre-commit checks

* Update current tests and add new test with proxy

airflow/providers/dbt/cloud/hooks/dbt.py

@@ -26,7 +26,6 @@
 from typing import TYPE_CHECKING, Any, Callable, Sequence, Set, TypeVar, cast
 
 import aiohttp
-from aiohttp import ClientResponseError
 from asgiref.sync import sync_to_async
 from requests.auth import AuthBase
 from requests.sessions import Session
@@ -182,9 +181,12 @@ class DbtCloudHook(HttpHook):
     def get_ui_field_behaviour(cls) -> dict[str, Any]:
         """Build custom field behavior for the dbt Cloud connection form in the Airflow UI."""
         return {
-            "hidden_fields": ["schema", "port", "extra"],
+            "hidden_fields": ["schema", "port"],
             "relabeling": {"login": "Account ID", "password": "API Token", "host": "Tenant"},
-            "placeholders": {"host": "Defaults to 'cloud.getdbt.com'."},
+            "placeholders": {
+                "host": "Defaults to 'cloud.getdbt.com'.",
+                "extra": "Optional JSON-formatted extra.",
+            },
         }
 
     def __init__(self, dbt_cloud_conn_id: str = default_conn_name, *args, **kwargs) -> None:
@@ -195,6 +197,10 @@ def __init__(self, dbt_cloud_conn_id: str = default_conn_name, *args, **kwargs)
     def _get_tenant_domain(conn: Connection) -> str:
         return conn.host or "cloud.getdbt.com"
 
+    @staticmethod
+    def _get_proxies(conn: Connection) -> dict[str, str] | None:
+        return conn.extra_dejson.get("proxies", None)
+
     @staticmethod
     def get_request_url_params(
         tenant: str, endpoint: str, include_related: list[str] | None = None, *, api_version: str = "v2"
@@ -238,14 +244,26 @@ async def get_job_details(
         endpoint = f"{account_id}/runs/{run_id}/"
         headers, tenant = await self.get_headers_tenants_from_connection()
         url, params = self.get_request_url_params(tenant, endpoint, include_related)
-        async with aiohttp.ClientSession(headers=headers) as session, session.get(
-            url, params=params
-        ) as response:
-            try:
-                response.raise_for_status()
-                return await response.json()
-            except ClientResponseError as e:
-                raise AirflowException(f"{e.status}:{e.message}")
+        proxies = self._get_proxies(self.connection)
+        async with aiohttp.ClientSession(headers=headers) as session:
+            if proxies is not None:
+                if url.startswith("https"):
+                    proxy = proxies.get("https")
+                else:
+                    proxy = proxies.get("http")
+                async with session.get(url, params=params, proxy=proxy) as response:
+                    try:
+                        response.raise_for_status()
+                        return await response.json()
+                    except aiohttp.ClientResponseError as e:
+                        raise AirflowException(f"{e.status}:{e.message}")
+            else:
+                async with session.get(url, params=params) as response:
+                    try:
+                        response.raise_for_status()
+                        return await response.json()
+                    except aiohttp.ClientResponseError as e:
+                        raise AirflowException(f"{e.status}:{e.message}")
 
     async def get_job_status(
         self, run_id: int, account_id: int | None = None, include_related: list[str] | None = None
@@ -280,8 +298,11 @@ def get_conn(self, *args, **kwargs) -> Session:
 
         return session
 
-    def _paginate(self, endpoint: str, payload: dict[str, Any] | None = None) -> list[Response]:
-        response = self.run(endpoint=endpoint, data=payload)
+    def _paginate(
+        self, endpoint: str, payload: dict[str, Any] | None = None, proxies: dict[str, str] | None = None
+    ) -> list[Response]:
+        extra_options = {"proxies": proxies} if proxies is not None else None
+        response = self.run(endpoint=endpoint, data=payload, extra_options=extra_options)
         resp_json = response.json()
         limit = resp_json["extra"]["filters"]["limit"]
         num_total_results = resp_json["extra"]["pagination"]["total_count"]
@@ -292,7 +313,7 @@ def _paginate(self, endpoint: str, payload: dict[str, Any] | None = None) -> lis
             _paginate_payload["offset"] = limit
 
             while num_current_results < num_total_results:
-                response = self.run(endpoint=endpoint, data=_paginate_payload)
+                response = self.run(endpoint=endpoint, data=_paginate_payload, extra_options=extra_options)
                 resp_json = response.json()
                 results.append(response)
                 num_current_results += resp_json["extra"]["pagination"]["count"]
@@ -310,17 +331,20 @@ def _run_and_get_response(
     ) -> Any:
         self.method = method
         full_endpoint = f"api/{api_version}/accounts/{endpoint}" if endpoint else None
+        proxies = self._get_proxies(self.connection)
+        extra_options = {"proxies": proxies} if proxies is not None else None
 
         if paginate:
             if isinstance(payload, str):
                 raise ValueError("Payload cannot be a string to paginate a response.")
 
             if full_endpoint:
-                return self._paginate(endpoint=full_endpoint, payload=payload)
+                return self._paginate(endpoint=full_endpoint, payload=payload, proxies=proxies)
 
             raise ValueError("An endpoint is needed to paginate a response.")
 
-        return self.run(endpoint=full_endpoint, data=payload)
+        # breakpoint()
+        return self.run(endpoint=full_endpoint, data=payload, extra_options=extra_options)
 
     def list_accounts(self) -> list[Response]:
         """

docs/apache-airflow-providers-dbt-cloud/connections.rst

@@ -77,6 +77,20 @@ Host (optional)
     If using the Connection form in the Airflow UI, the Tenant domain can also be stored in the "Tenant"
     field.
 
+Extra (optional)
+    Specify extra parameters as JSON dictionary. As of now, only `proxies` is supported when wanting to connect to dbt Cloud via a proxy.
+
+    `proxies` should be a dictionary of proxies to be used by HTTP and HTTPS connections.
+
+.. code-block:: json
+
+    {
+      "proxies": {
+        "http": "http://myproxy.mycompany.local:8080",
+        "https": "http://myproxy.mycompany.local:8080"
+      }
+    }
+
 When specifying the connection as an environment variable, you should specify it following the standard syntax
 of a database connection. Note that all components of the URI should be URL-encoded.