diff --git a/providers/apache/hive/src/airflow/providers/apache/hive/hooks/hive.py b/providers/apache/hive/src/airflow/providers/apache/hive/hooks/hive.py index b72f87f87162c..7ba8834694fff 100644 --- a/providers/apache/hive/src/airflow/providers/apache/hive/hooks/hive.py +++ b/providers/apache/hive/src/airflow/providers/apache/hive/hooks/hive.py @@ -865,7 +865,7 @@ def get_conn(self, schema: str | None = None) -> Any: username: str | None = None password: str | None = None - db = self.get_connection(self.hiveserver2_conn_id) # type: ignore + db = self.get_connection(self.get_conn_id()) auth_mechanism = db.extra_dejson.get("auth_mechanism", "NONE") if auth_mechanism == "NONE" and db.login is None: @@ -907,7 +907,7 @@ def _get_results( with contextlib.closing(self.get_conn(schema)) as conn, contextlib.closing(conn.cursor()) as cur: cur.arraysize = fetch_size or 1000 - db = self.get_connection(self.hiveserver2_conn_id) # type: ignore + db = self.get_connection(self.get_conn_id()) # Not all query services (e.g. impala) support the set command if db.extra_dejson.get("run_set_variable_statements", True): env_context = get_context_from_env_var() @@ -1029,9 +1029,10 @@ def get_records( schema = kwargs["schema"] if "schema" in kwargs else "default" return self.get_results(sql, schema=schema, hive_conf=parameters)["data"] - def _get_pandas_df( # type: ignore + def _get_pandas_df( self, - sql: str, + sql, + parameters: list[Any] | tuple[Any, ...] | Mapping[str, Any] | None = None, schema: str = "default", hive_conf: dict[Any, Any] | None = None, **kwargs, @@ -1047,9 +1048,10 @@ def _get_pandas_df( # type: ignore df = pd.DataFrame(res["data"], columns=[c[0] for c in res["header"]], **kwargs) return df - def _get_polars_df( # type: ignore + def _get_polars_df( self, - sql: str, + sql, + parameters: list[Any] | tuple[Any, ...] | Mapping[str, Any] | None = None, schema: str = "default", hive_conf: dict[Any, Any] | None = None, **kwargs, @@ -1076,7 +1078,7 @@ def get_df( **kwargs: Any, ) -> pd.DataFrame: ... - @overload # type: ignore[override] + @overload def get_df( self, sql: str, @@ -1087,7 +1089,7 @@ def get_df( **kwargs: Any, ) -> pl.DataFrame: ... - def get_df( # type: ignore + def get_df( self, sql: str, schema: str = "default",