diff --git a/src/firebolt/async_db/connection.py b/src/firebolt/async_db/connection.py index 382b19d20a1..c7d2dc21fa2 100644 --- a/src/firebolt/async_db/connection.py +++ b/src/firebolt/async_db/connection.py @@ -4,10 +4,10 @@ import socket from json import JSONDecodeError from types import TracebackType -from typing import Any, Dict, Iterable, List, Optional +from typing import Any, Dict, List, Optional from httpcore.backends.auto import AutoBackend -from httpcore.backends.base import SOCKET_OPTION, AsyncNetworkStream +from httpcore.backends.base import AsyncNetworkStream from httpx import AsyncHTTPTransport, HTTPStatusError, RequestError, Timeout from firebolt.async_db.cursor import Cursor @@ -134,16 +134,20 @@ class OverriddenHttpBackend(AutoBackend): and `KEEPIDLE` settings. """ - async def connect_tcp( + async def connect_tcp( # type: ignore [override] self, host: str, port: int, timeout: Optional[float] = None, local_address: Optional[str] = None, - socket_options: Optional[Iterable[SOCKET_OPTION]] = None, + **kwargs: Any, ) -> AsyncNetworkStream: - stream = await super().connect_tcp( - host, port, timeout=timeout, local_address=local_address + stream = await super().connect_tcp( # type: ignore [call-arg] + host, + port, + timeout=timeout, + local_address=local_address, + **kwargs, ) # Enable keepalive stream.get_extra_info("socket").setsockopt( diff --git a/src/firebolt/db/connection.py b/src/firebolt/db/connection.py index e4b2a74e088..3a29d7852e5 100644 --- a/src/firebolt/db/connection.py +++ b/src/firebolt/db/connection.py @@ -4,10 +4,10 @@ import socket from json import JSONDecodeError from types import TracebackType -from typing import Any, Dict, Iterable, List, Optional +from typing import Any, Dict, List, Optional from warnings import warn -from httpcore.backends.base import SOCKET_OPTION, NetworkStream +from httpcore.backends.base import NetworkStream from httpcore.backends.sync import SyncBackend from httpx import HTTPStatusError, HTTPTransport, RequestError, Timeout from readerwriterlock.rwlock import RWLockWrite @@ -125,20 +125,20 @@ class OverriddenHttpBackend(SyncBackend): and `KEEPIDLE` settings. """ - def connect_tcp( + def connect_tcp( # type: ignore [override] self, host: str, port: int, timeout: Optional[float] = None, local_address: Optional[str] = None, - socket_options: Optional[Iterable[SOCKET_OPTION]] = None, + **kwargs: Any, ) -> NetworkStream: - stream = super().connect_tcp( + stream = super().connect_tcp( # type: ignore [call-arg] host, port, timeout=timeout, local_address=local_address, - socket_options=socket_options, + **kwargs, ) # Enable keepalive stream.get_extra_info("socket").setsockopt(