Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
76 changes: 76 additions & 0 deletions task-sdk/src/airflow/sdk/definitions/connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -290,6 +290,82 @@ async def async_get(cls, conn_id: str) -> Any:
except AirflowRuntimeError as e:
cls._handle_connection_error(e, conn_id)

@property
async def aextra_dejson(self) -> dict:
"""Async version: Returns the extra property by deserializing json and masking secrets."""
from airflow.sdk.log import amask_secret

extra = {}
if self.extra:
try:
import json
extra = json.loads(self.extra)
except Exception:
log.exception("Failed to deserialize extra property `extra`, returning empty dictionary")
else:
await amask_secret(extra)
return extra

async def aget_uri(self) -> str:
"""Async version: Generate and return connection in URI format."""
from urllib.parse import parse_qsl, quote, urlencode

if self.conn_type:
uri = f"{self.conn_type.lower().replace('_', '-')}://"
else:
uri = "//"
host_to_use: str | None
if self.host and "://" in self.host:
protocol, host = self.host.split("://", 1)
# If the protocol in host matches the connection type, don't add it again
if protocol == self.conn_type:
host_to_use = self.host
protocol_to_add = None
else:
# Different protocol, add it to the URI
host_to_use = host
protocol_to_add = protocol
else:
host_to_use = self.host
protocol_to_add = None

if protocol_to_add:
uri += f"{protocol_to_add}://"

authority_block = ""
if self.login is not None:
authority_block += quote(self.login, safe="")
if self.password is not None:
authority_block += ":" + quote(self.password, safe="")
if authority_block > "":
authority_block += "@"
uri += authority_block

host_block = ""
if host_to_use:
host_block += quote(host_to_use, safe="")
if self.port:
if host_block == "" and authority_block == "":
host_block += f"@:{self.port}"
else:
host_block += f":{self.port}"
if self.schema:
host_block += f"/{quote(self.schema, safe='')}"
uri += host_block

if self.extra:
try:
extra_dejson = await self.aextra_dejson()
query: str | None = urlencode(extra_dejson)
except TypeError:
query = None
if query and extra_dejson == dict(parse_qsl(query, keep_blank_values=True)):
uri += ("?" if self.schema else "/?") + query
else:
uri += ("?" if self.schema else "/?") + urlencode({self.EXTRA_KEY: self.extra})

return uri

@property
def extra_dejson(self) -> dict:
"""Returns the extra property by deserializing json."""
Expand Down
2 changes: 1 addition & 1 deletion task-sdk/src/airflow/sdk/execution_time/context.py
Original file line number Diff line number Diff line change
Expand Up @@ -297,7 +297,7 @@ async def _async_get_connection(conn_id: str) -> Connection:
conn = await sync_to_async(secrets_backend.get_connection)(conn_id) # type: ignore[assignment]

if conn:
SecretCache.save_connection_uri(conn_id, conn.get_uri())
SecretCache.save_connection_uri(conn_id, await conn.aget_uri())
await _amask_connection_secrets(conn)
return conn
except AirflowSecretsBackendAccessDenied:
Expand Down
Loading