from typing import Any, Dict, Optional
from airflow.exceptions import AirflowException
from airflow.utils.decorators import apply_defaults
from config import Config
from operators.superset.add_superset_db_connection import (
AddSupersetDbConnectionOperator,
)
[docs]
class AddSupersetDbConnectionFromSecretOperator(AddSupersetDbConnectionOperator):
"""Create a Superset DB connection using an AWS Secrets Manager secret.
The SQLAlchemy URI is built inside ``execute()`` so the password never
appears in XCom or task return values.
Expected secret JSON fields:
username, password, host, database (or ``dbname``).
port - optional, defaults to 5439.
:param database_name: Display label shown in the Superset UI. This is
NOT the actual database name connected to — the real DB name comes
from the secret's ``database`` (or ``dbname``) field. Free-form;
spaces allowed.
:type database_name: str
:param secret_name: AWS Secrets Manager secret id or ARN.
:type secret_name: str
:param dialect: SQLAlchemy dialect prefix (default ``redshift+psycopg2``).
:type dialect: str
:param default_port: Port to use when the secret omits one. Defaults to 5439.
:type default_port: int
:param expose_in_sqllab: Whether to expose the database in SQL Lab.
:type expose_in_sqllab: bool
:param allow_ctas: Whether to allow CTAS statements.
:type allow_ctas: bool
:param allow_cvas: Whether to allow CVAS statements.
:type allow_cvas: bool
:param allow_dml: Whether to allow DML statements.
:type allow_dml: bool
:param extra: Optional dict merged into the Superset payload.
:type extra: dict
:param conn_id: Airflow connection id for the Superset instance.
:type conn_id: str
:param timeout: HTTP request timeout in seconds.
:type timeout: int
"""
template_fields = (
"database_name",
"secret_name",
"extra",
"conn_id",
"base_url",
"owner_email",
)
ui_color = "#fce4ec"
@apply_defaults
def __init__(
self,
database_name: str,
secret_name: str,
dialect: str = "redshift+psycopg2",
default_port: int = 5439,
expose_in_sqllab: bool = True,
allow_ctas: bool = False,
allow_cvas: bool = False,
allow_dml: bool = False,
extra: Optional[Dict[str, Any]] = None,
conn_id: str = Config.SUPERSET_DAP_CONN_ID,
base_url: Optional[str] = None,
timeout: int = 30,
owner_email: Optional[str] = None,
*args,
**kwargs,
):
super().__init__(
database_name=database_name,
sqlalchemy_uri="",
expose_in_sqllab=expose_in_sqllab,
allow_ctas=allow_ctas,
allow_cvas=allow_cvas,
allow_dml=allow_dml,
extra=extra,
conn_id=conn_id,
base_url=base_url,
timeout=timeout,
owner_email=owner_email,
*args,
**kwargs,
)
self.secret_name = secret_name
self.dialect = dialect
self.default_port = default_port
[docs]
def execute(self, context) -> Dict[str, Any]:
import json
from urllib.parse import quote_plus
from util import Util
if not self.secret_name:
raise AirflowException("Missing required field: secret_name")
raw = Util.get_aws_secret(self.secret_name)
try:
secret = json.loads(raw)
except (TypeError, json.JSONDecodeError) as e:
raise AirflowException(
f"Secret {self.secret_name} is not valid JSON: {e}"
)
database = secret.get("database") or secret.get("dbname")
host = secret.get("host")
fields = {
"username": secret.get("username"),
"password": secret.get("password"),
"host": host,
"database": database,
}
missing = [k for k, v in fields.items() if not v]
if missing:
raise AirflowException(
f"Secret {self.secret_name} missing required keys: "
f"{', '.join(missing)}"
)
port = secret.get("port", self.default_port)
# Percent-encode userinfo so URI-reserved characters in credentials
# (@, #, %, :, /, ?, &, +, space, ...) don't break the SQLAlchemy URI.
# SQLAlchemy reliably decodes the userinfo segment back to the
# original values. The host and database segments are left unencoded
# because DNS-valid hosts and conventional DB names contain no
# reserved characters, and SQLAlchemy's decoding of the path segment
# varies by dialect.
user = quote_plus(str(fields["username"]))
pwd = quote_plus(str(fields["password"]))
host = fields["host"]
self.sqlalchemy_uri = (
f"{self.dialect}://{user}:{pwd}@{host}:{port}/{fields['database']}"
)
self.log.info(
"Built SQLAlchemy URI for host=%s port=%s db=%s "
"(credentials redacted)",
host,
port,
fields["database"],
)
return super().execute(context)