Skip to content

[PECO-1137] Reintroduce protocol checking to Python test fw #248

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
nithinkdb merged 10 commits into from
Oct 25, 2023
Merged

[PECO-1137] Reintroduce protocol checking to Python test fw #248

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
67a0e00
Put in some unit tests, will add e2e
nithinkdb Oct 5, 2023
4841759
Added e2e test
nithinkdb Oct 16, 2023
9763274
Linted
nithinkdb Oct 16, 2023
d3d36ea
re-bumped thrift files
nithinkdb Oct 16, 2023
74a1531
Changed structure to store protocol version as feature of connection
nithinkdb Oct 16, 2023
b06791b
Fixed parameters test
nithinkdb Oct 24, 2023
e24b0b3
Fixed comments
nithinkdb Oct 24, 2023
af468e9
Update src/databricks/sql/client.py
nithinkdb Oct 25, 2023
d8b046e
Fixed comments
nithinkdb Oct 25, 2023
05ea379
Removed extra indent
nithinkdb Oct 25, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 35 additions & 1 deletion src/databricks/sql/client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
SessionAlreadyClosedError,
CursorAlreadyClosedError,
)
from databricks.sql.thrift_api.TCLIService import ttypes
from databricks.sql.thrift_backend import ThriftBackend
from databricks.sql.utils import (
ExecuteResponse,
Expand Down Expand Up @@ -196,9 +197,11 @@ def read(self) -> Optional[OAuthToken]:
**kwargs,
)

self._session_handle = self.thrift_backend.open_session(
self._open_session_resp = self.thrift_backend.open_session(
session_configuration, catalog, schema
)
self._session_handle = self._open_session_resp.sessionHandle
self.protocol_version = self.get_protocol_version(self._open_session_resp)
self.use_cloud_fetch = kwargs.get("use_cloud_fetch", True)
self.open = True
logger.info("Successfully opened session " + str(self.get_session_id_hex()))
Expand All @@ -225,6 +228,30 @@ def __del__(self):
def get_session_id(self):
return self.thrift_backend.handle_to_id(self._session_handle)

@staticmethod
def get_protocol_version(openSessionResp):
"""
Since the sessionHandle will sometimes have a serverProtocolVersion, it takes
precedence over the serverProtocolVersion defined in the OpenSessionResponse.
"""
if (
openSessionResp.sessionHandle
and hasattr(openSessionResp.sessionHandle, "serverProtocolVersion")
and openSessionResp.sessionHandle.serverProtocolVersion
):
return openSessionResp.sessionHandle.serverProtocolVersion
return openSessionResp.serverProtocolVersion

@staticmethod
def server_parameterized_queries_enabled(protocolVersion):
if (
protocolVersion
and protocolVersion >= ttypes.TProtocolVersion.SPARK_CLI_SERVICE_PROTOCOL_V8
):
return True
else:
return False

def get_session_id_hex(self):
return self.thrift_backend.handle_to_hex_id(self._session_handle)

Expand Down Expand Up @@ -501,6 +528,13 @@ def execute(
"""
if parameters is None:
parameters = []

elif not Connection.server_parameterized_queries_enabled(
self.connection.protocol_version
):
raise NotSupportedError(
"Parameterized operations are not supported by this server. DBR 14.1 is required."
)
else:
parameters = named_parameters_to_tsparkparams(parameters)

Expand Down
Loading