diff --git a/providers/trino/src/airflow/providers/trino/hooks/trino.py b/providers/trino/src/airflow/providers/trino/hooks/trino.py index 35bb83cce2a8a..40d7c8530427a 100644 --- a/providers/trino/src/airflow/providers/trino/hooks/trino.py +++ b/providers/trino/src/airflow/providers/trino/hooks/trino.py @@ -146,14 +146,14 @@ def __init__(self, *args, **kwargs): def get_conn(self) -> Connection: """Return a connection object.""" - db = self.get_connection(self.trino_conn_id) # type: ignore[attr-defined] + db = self.get_connection(self.get_conn_id()) extra = db.extra_dejson auth = None user = db.login if db.password and extra.get("auth") in ("kerberos", "certs"): raise AirflowException(f"The {extra.get('auth')!r} authorization type doesn't support password.") if db.password: - auth = trino.auth.BasicAuthentication(db.login, db.password) # type: ignore[attr-defined] + auth = trino.auth.BasicAuthentication(db.login, db.password) elif extra.get("auth") == "jwt": if not exactly_one(jwt_file := "jwt__file" in extra, jwt_token := "jwt__token" in extra): msg = ( @@ -176,7 +176,7 @@ def get_conn(self) -> Connection: extra.get("certs__client_key_path"), ) elif extra.get("auth") == "kerberos": - auth = trino.auth.KerberosAuthentication( # type: ignore[attr-defined] + auth = trino.auth.KerberosAuthentication( config=extra.get("kerberos__config", os.environ.get("KRB5_CONFIG")), service_name=extra.get("kerberos__service_name"), mutual_authentication=_boolify(extra.get("kerberos__mutual_authentication", False)), @@ -205,7 +205,6 @@ def get_conn(self) -> Connection: catalog=extra.get("catalog", "hive"), schema=db.schema, auth=auth, - # type: ignore[func-returns-value] isolation_level=self.get_isolation_level(), verify=_boolify(extra.get("verify", True)), session_properties=extra.get("session_properties") or None, @@ -219,7 +218,7 @@ def get_conn(self) -> Connection: def get_isolation_level(self) -> Any: """Return an isolation level.""" - db = self.get_connection(self.trino_conn_id) # type: ignore[attr-defined] + db = self.get_connection(self.get_conn_id()) isolation_level = db.extra_dejson.get("isolation_level", "AUTOCOMMIT").upper() return getattr(IsolationLevel, isolation_level, IsolationLevel.AUTOCOMMIT)