Skip to content

Commit

Permalink
Satisfy Mypy for Livy Hook (#46157)
Browse files Browse the repository at this point in the history
The Livy Hook shows ... strange MyPy errors and hard to say why.

First of all it incorrectly used multiple inheritance (BaseHook
already inherist from LoggingMixin) but then it seems that mypy
has been lost when it comes to fields available in the hook.
Setting them directly (like it happens in other hooks deriving
from HttpHook solves the MyPy issues.
  • Loading branch information
potiuk authored Jan 28, 2025
1 parent bc5e506 commit 031028d
Showing 1 changed file with 8 additions and 5 deletions.
13 changes: 8 additions & 5 deletions providers/src/airflow/providers/apache/livy/hooks/livy.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,6 @@

from airflow.exceptions import AirflowException
from airflow.providers.http.hooks.http import HttpAsyncHook, HttpHook
from airflow.utils.log.logging_mixin import LoggingMixin

if TYPE_CHECKING:
from airflow.models import Connection
Expand All @@ -53,7 +52,7 @@ class BatchState(Enum):
SUCCESS = "success"


class LivyHook(HttpHook, LoggingMixin):
class LivyHook(HttpHook):
"""
Hook for Apache Livy through the REST API.
Expand Down Expand Up @@ -88,7 +87,9 @@ def __init__(
extra_headers: dict[str, Any] | None = None,
auth_type: Any | None = None,
) -> None:
super().__init__(http_conn_id=livy_conn_id)
super().__init__()
self.method = "POST"
self.http_conn_id = livy_conn_id
self.extra_headers = extra_headers or {}
self.extra_options = extra_options or {}
if auth_type:
Expand Down Expand Up @@ -457,7 +458,7 @@ def _validate_extra_conf(conf: dict[Any, Any]) -> bool:
return True


class LivyAsyncHook(HttpAsyncHook, LoggingMixin):
class LivyAsyncHook(HttpAsyncHook):
"""
Hook for Apache Livy through the REST API asynchronously.
Expand Down Expand Up @@ -490,7 +491,9 @@ def __init__(
extra_options: dict[str, Any] | None = None,
extra_headers: dict[str, Any] | None = None,
) -> None:
super().__init__(http_conn_id=livy_conn_id)
super().__init__()
self.method = "POST"
self.http_conn_id = livy_conn_id
self.extra_headers = extra_headers or {}
self.extra_options = extra_options or {}

Expand Down

0 comments on commit 031028d

Please sign in to comment.