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.