Repository: apache/airflow
File: providers/amazon/src/airflow/providers/amazon/aws/hooks/athena_sql.py (Python)

Code before:
    Also, aws_domain could be specified in ``extra`` of your connection.

    PyAthena and AWS Authentication parameters could be passed in extra field of ``athena_conn_id`` connection.

    Passing authentication parameters in ``athena_conn_id`` will override those in ``aws_conn_id``.

    .. seealso::
        :class:`~airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook`

    .. note::
        get_uri() depends on SQLAlchemy and PyAthena
    """

    conn_name_attr = "athena_conn_id"
    default_conn_name = "athena_default"
    conn_type = "athena"
    hook_name = "Amazon Athena"
    supports_autocommit = True

    def __init__(self, athena_conn_id: str = default_conn_name, *args, **kwargs) -> None:
        # AwsGenericHook.__init__() only accepts these kwargs. Connection extras
        # like s3_staging_dir and work_group are not constructor params — they are
        # read later from the connection in get_conn(). BaseSQLOperator.get_hook()
        # passes all connection extras as kwargs, so we must filter them out here.
        _aws_generic_hook_kwargs = {"aws_conn_id", "verify", "region_name", "client_type", "resource_type", "config"}
        filtered_kwargs = {k: v for k, v in kwargs.items() if k in _aws_generic_hook_kwargs}
        super().__init__(*args, **filtered_kwargs)
        self.athena_conn_id = athena_conn_id

    @classmethod
    def get_ui_field_behaviour(cls) -> dict[str, Any]:
        """Return custom UI field behaviour for AWS Athena Connection."""
        return {
            "hidden_fields": ["host", "port"],
            "relabeling": {
                "login": "AWS Access Key ID",
                "password": "AWS Secret Access Key",
            },
            "placeholders": {
                "login": "AKIAIOSFODNN7EXAMPLE",
                "password": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
                "extra": json.dumps(
                    {
                        "aws_domain": "amazonaws.com",
                        "driver": "rest",
                        "s3_staging_dir": "s3://bucket_name/staging/",
                        "work_group": "primary",
                        "region_name": "us-east-1",
                        "session_kwargs": {"profile_name": "default"},
                        "config_kwargs": {"retries": {"mode": "standard", "max_attempts": 10}},
                        "role_arn": "arn:aws:iam::123456789098:role/role-name",

Review comment (suggestion):
Instead of doing that, can you declare these parameters in the constructor? That would avoid doing a filtering

Code after:

    Also, aws_domain could be specified in ``extra`` of your connection.

    PyAthena and AWS Authentication parameters could be passed in extra field of ``athena_conn_id`` connection.

    Passing authentication parameters in ``athena_conn_id`` will override those in ``aws_conn_id``.

    .. seealso::
        :class:`~airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook`

    .. note::
        get_uri() depends on SQLAlchemy and PyAthena
    """

    conn_name_attr = "athena_conn_id"
    default_conn_name = "athena_default"
    conn_type = "athena"
    hook_name = "Amazon Athena"
    supports_autocommit = True

    def __init__(
        self,
        athena_conn_id: str = default_conn_name,
        aws_conn_id: str | None = AwsBaseHook.default_conn_name,
        verify: bool | str | None = None,
        region_name: str | None = None,
        client_type: str | None = None,
        resource_type: str | None = None,
        config: Config | dict[str, Any] | None = None,
        **kwargs,
    ) -> None:
        # AwsGenericHook.__init__() only accepts the params declared above.
        # Connection extras like s3_staging_dir and work_group are not
        # constructor params — they are read later from the connection in
        # get_conn(). BaseSQLOperator.get_hook() passes all connection extras
        # as kwargs, so we absorb them here via **kwargs and discard them.
        super().__init__(
            aws_conn_id=aws_conn_id,
            verify=verify,
            region_name=region_name,
            client_type=client_type,
            resource_type=resource_type,
            config=config,
        )
        self.athena_conn_id = athena_conn_id

    @classmethod
    def get_ui_field_behaviour(cls) -> dict[str, Any]:
        """Return custom UI field behaviour for AWS Athena Connection."""
        return {
            "hidden_fields": ["host", "port"],