Skip to content
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -1076,7 +1078,7 @@ def get_df(
**kwargs: Any,
) -> pd.DataFrame: ...

@overload # type: ignore[override]
@overload
def get_df(
self,
sql: str,
Expand All @@ -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",
Expand Down