Jul 10, 2025 · Aug 17, 2022 · Aug 17, 2022 · Aug 19, 2022 · Aug 19, 2022 · Aug 23, 2022
diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py
 import time
 from typing import Dict, Tuple, List, Optional, Any, Union, Sequence

 import pandas

 try:
    OperationalError,
    SessionAlreadyClosedError,
    CursorAlreadyClosedError,
    InterfaceError,
    NotSupportedError,
    ProgrammingError,
 )
 from databricks.sql.thrift_api.TCLIService import ttypes
 from databricks.sql.thrift_backend import ThriftBackend
    TSparkParameter,
    TOperationState,
 )

 from databricks.sql.telemetry.telemetry_client import (
    TelemetryHelper,
    TelemetryClientFactory,
 )
 from databricks.sql.telemetry.models.enums import DatabricksClientType
 from databricks.sql.telemetry.models.event import (
    DriverConnectionParameters,
    HostDetails,
 )
 from databricks.sql.telemetry.latency_logger import log_latency
 from databricks.sql.telemetry.models.enums import StatementType

 logger = logging.getLogger(__name__)

            server_hostname, **kwargs
        )

        self.server_telemetry_enabled = True
        self.client_telemetry_enabled = kwargs.get("enable_telemetry", False)
        self.telemetry_enabled = (
            self.client_telemetry_enabled and self.server_telemetry_enabled
        )

        user_agent_entry = kwargs.get("user_agent_entry")
        if user_agent_entry is None:
            user_agent_entry = kwargs.get("_user_agent_entry")
            kwargs.get("use_inline_params", False)
        )

        TelemetryClientFactory.initialize_telemetry_client(
            telemetry_enabled=self.telemetry_enabled,
            session_id_hex=self.get_session_id_hex(),
            auth_provider=auth_provider,
            host_url=self.host,
        )

        self._telemetry_client = TelemetryClientFactory.get_telemetry_client(
            session_id_hex=self.get_session_id_hex()
        )

        driver_connection_params = DriverConnectionParameters(
            http_path=http_path,
            mode=DatabricksClientType.THRIFT,
            host_info=HostDetails(host_url=server_hostname, port=self.port),
            auth_mech=TelemetryHelper.get_auth_mechanism(auth_provider),
            auth_flow=TelemetryHelper.get_auth_flow(auth_provider),
            socket_timeout=kwargs.get("_socket_timeout", None),
        )

        self._telemetry_client.export_initial_telemetry_log(
            driver_connection_params=driver_connection_params,
            user_agent=useragent_header,
        )

    def _set_use_inline_params_with_warning(self, value: Union[bool, str]):
        """Valid values are True, False, and "silent"

        Will throw an Error if the connection has been closed.
        """
        if not self.open:
            raise Error("Cannot create cursor from closed connection")
            raise InterfaceError(
                "Cannot create cursor from closed connection",
                session_id_hex=self.get_session_id_hex(),
            )

        cursor = Cursor(
            self,

        self.open = False

        TelemetryClientFactory.close(self.get_session_id_hex())

    def commit(self):
        """No-op because Databricks does not support transactions"""
        pass

    def rollback(self):
        raise NotSupportedError("Transactions are not supported on Databricks")
        raise NotSupportedError(
            "Transactions are not supported on Databricks",
            session_id_hex=self.get_session_id_hex(),
        )


 class Cursor:
            for row in self.active_result_set:
                yield row
        else:
            raise Error("There is no active result set")
            raise ProgrammingError(
                "There is no active result set",
                session_id_hex=self.connection.get_session_id_hex(),
            )

    def _determine_parameter_approach(
        self, params: Optional[TParameterCollection]

    def _check_not_closed(self):
        if not self.open:
            raise Error("Attempting operation on closed cursor")
            raise InterfaceError(
                "Attempting operation on closed cursor",
                session_id_hex=self.connection.get_session_id_hex(),
            )

    def _handle_staging_operation(
        self, staging_allowed_local_path: Union[None, str, List[str]]
        elif isinstance(staging_allowed_local_path, type(list())):
            _staging_allowed_local_paths = staging_allowed_local_path
        else:
            raise Error(
                "You must provide at least one staging_allowed_local_path when initialising a connection to perform ingestion commands"
            raise ProgrammingError(
                "You must provide at least one staging_allowed_local_path when initialising a connection to perform ingestion commands",
                session_id_hex=self.connection.get_session_id_hex(),
            )

        abs_staging_allowed_local_paths = [
                else:
                    continue
            if not allow_operation:
                raise Error(
                    "Local file operations are restricted to paths within the configured staging_allowed_local_path"
                raise ProgrammingError(
                    "Local file operations are restricted to paths within the configured staging_allowed_local_path",
                    session_id_hex=self.connection.get_session_id_hex(),
                )

        # May be real headers, or could be json string
            handler_args.pop("local_file")
            return self._handle_staging_remove(**handler_args)
        else:
            raiseError(
            raiseProgrammingError(
                f"Operation {row.operation} is not supported. "
                + "Supported operations are GET, PUT, and REMOVE"
                + "Supported operations are GET, PUT, and REMOVE",
                session_id_hex=self.connection.get_session_id_hex(),
            )

    @log_latency(StatementType.SQL)
    def _handle_staging_put(
        self, presigned_url: str, local_file: str, headers: Optional[dict] = None
    ):
        """

        if local_file is None:
            raise Error("Cannot perform PUT without specifying a local_file")
            raise ProgrammingError(
                "Cannot perform PUT without specifying a local_file",
                session_id_hex=self.connection.get_session_id_hex(),
            )

        with open(local_file, "rb") as fh:
            r = requests.put(url=presigned_url, data=fh, headers=headers)
        # fmt: on

        if r.status_code not in [OK, CREATED, NO_CONTENT, ACCEPTED]:
            raise Error(
                f"Staging operation over HTTP was unsuccessful: {r.status_code}-{r.text}"
            raise OperationalError(
                f"Staging operation over HTTP was unsuccessful: {r.status_code}-{r.text}",
                session_id_hex=self.connection.get_session_id_hex(),
            )

        if r.status_code == ACCEPTED:
                + "but not yet applied on the server. It's possible this command may fail later."
            )

    @log_latency(StatementType.SQL)
    def _handle_staging_get(
        self, local_file: str, presigned_url: str, headers: Optional[dict] = None
    ):
        """

        if local_file is None:
            raise Error("Cannot perform GET without specifying a local_file")
            raise ProgrammingError(
                "Cannot perform GET without specifying a local_file",
                session_id_hex=self.connection.get_session_id_hex(),
            )

        r = requests.get(url=presigned_url, headers=headers)

        # response.ok verifies the status code is not between 400-600.
        # Any 2xx or 3xx will evaluate r.ok == True
        if not r.ok:
            raise Error(
                f"Staging operation over HTTP was unsuccessful: {r.status_code}-{r.text}"
            raise OperationalError(
                f"Staging operation over HTTP was unsuccessful: {r.status_code}-{r.text}",
                session_id_hex=self.connection.get_session_id_hex(),
            )

        with open(local_file, "wb") as fp:
            fp.write(r.content)

    @log_latency(StatementType.SQL)
    def _handle_staging_remove(
        self, presigned_url: str, headers: Optional[dict] = None
    ):
        r = requests.delete(url=presigned_url, headers=headers)

        if not r.ok:
            raise Error(
                f"Staging operation over HTTP was unsuccessful: {r.status_code}-{r.text}"
            raise OperationalError(
                f"Staging operation over HTTP was unsuccessful: {r.status_code}-{r.text}",
                session_id_hex=self.connection.get_session_id_hex(),
            )

    @log_latency(StatementType.QUERY)
    def execute(
        self,
        operation: str,

        return self

    @log_latency(StatementType.QUERY)
    def execute_async(
        self,
        operation: str,

            return self
        else:
            raise Error(
                f"get_execution_result failed with Operation status {operation_state}"
            raise OperationalError(
                f"get_execution_result failed with Operation status {operation_state}",
                session_id_hex=self.connection.get_session_id_hex(),
            )

    def executemany(self, operation, seq_of_parameters):
            self.execute(operation, parameters)
        return self

    @log_latency(StatementType.METADATA)
    def catalogs(self) -> "Cursor":
        """
        Get all available catalogs.
        )
        return self

    @log_latency(StatementType.METADATA)
    def schemas(
        self, catalog_name: Optional[str] = None, schema_name: Optional[str] = None
    ) -> "Cursor":
        )
        return self

    @log_latency(StatementType.METADATA)
    def tables(
        self,
        catalog_name: Optional[str] = None,
        )
        return self

    @log_latency(StatementType.METADATA)
    def columns(
        self,
        catalog_name: Optional[str] = None,
        if self.active_result_set:
            return self.active_result_set.fetchall()
        else:
            raise Error("There is no active result set")
            raise ProgrammingError(
                "There is no active result set",
                session_id_hex=self.connection.get_session_id_hex(),
            )

    def fetchone(self) -> Optional[Row]:
        """
        if self.active_result_set:
            return self.active_result_set.fetchone()
        else:
            raise Error("There is no active result set")
            raise ProgrammingError(
                "There is no active result set",
                session_id_hex=self.connection.get_session_id_hex(),
            )

    def fetchmany(self, size: int) -> List[Row]:
        """
        if self.active_result_set:
            return self.active_result_set.fetchmany(size)
        else:
            raise Error("There is no active result set")
            raise ProgrammingError(
                "There is no active result set",
                session_id_hex=self.connection.get_session_id_hex(),
            )

    def fetchall_arrow(self) -> "pyarrow.Table":
        self._check_not_closed()
        if self.active_result_set:
            return self.active_result_set.fetchall_arrow()
        else:
            raise Error("There is no active result set")
            raise ProgrammingError(
                "There is no active result set",
                session_id_hex=self.connection.get_session_id_hex(),
            )

    def fetchmany_arrow(self, size) -> "pyarrow.Table":
        self._check_not_closed()
        if self.active_result_set:
            return self.active_result_set.fetchmany_arrow(size)
        else:
            raise Error("There is no active result set")
            raise ProgrammingError(
                "There is no active result set",
                session_id_hex=self.connection.get_session_id_hex(),
            )

    def cancel(self) -> None:
        """

        return results

    @log_latency()
    def fetchone(self) -> Optional[Row]:
        """
        Fetch the next row of a query result set, returning a single sequence,
        else:
            return None

    @log_latency()
    def fetchall(self) -> List[Row]:
        """
        Fetch all (remaining) rows of a query result, returning them as a list of rows.
        else:
            return self._convert_arrow_table(self.fetchall_arrow())

    @log_latency()
    def fetchmany(self, size: int) -> List[Row]:
        """
        Fetch the next set of rows of a query result, returning a list of rows.
Original file line number	Diff line number	Diff line change
		@@ -1,6 +1,5 @@
		import time
		from typing import Dict, Tuple, List, Optional, Any, Union, Sequence

		import pandas

		try:
Expand All		@@ -19,6 +18,9 @@
		OperationalError,
		SessionAlreadyClosedError,
		CursorAlreadyClosedError,
		InterfaceError,
		NotSupportedError,
		ProgrammingError,
		)
		from databricks.sql.thrift_api.TCLIService import ttypes
		from databricks.sql.thrift_backend import ThriftBackend
Expand DownExpand Up		@@ -50,7 +52,17 @@
		TSparkParameter,
		TOperationState,
		)

		from databricks.sql.telemetry.telemetry_client import (
		TelemetryHelper,
		TelemetryClientFactory,
		)
		from databricks.sql.telemetry.models.enums import DatabricksClientType
		from databricks.sql.telemetry.models.event import (
		DriverConnectionParameters,
		HostDetails,
		)
		from databricks.sql.telemetry.latency_logger import log_latency
		from databricks.sql.telemetry.models.enums import StatementType

		logger = logging.getLogger(__name__)

Expand DownExpand Up		@@ -234,6 +246,12 @@ def read(self) -> Optional[OAuthToken]:
		server_hostname, **kwargs
		)

		self.server_telemetry_enabled = True
		self.client_telemetry_enabled = kwargs.get("enable_telemetry", False)
saishreeeee marked this conversation as resolved. Show resolvedHide resolved
		self.telemetry_enabled = (
		self.client_telemetry_enabled and self.server_telemetry_enabled
		)

		user_agent_entry = kwargs.get("user_agent_entry")
		if user_agent_entry is None:
		user_agent_entry = kwargs.get("_user_agent_entry")
Expand DownExpand Up		@@ -289,6 +307,31 @@ def read(self) -> Optional[OAuthToken]:
		kwargs.get("use_inline_params", False)
		)

		TelemetryClientFactory.initialize_telemetry_client(
		telemetry_enabled=self.telemetry_enabled,
		session_id_hex=self.get_session_id_hex(),
		auth_provider=auth_provider,
		host_url=self.host,
		)

		self._telemetry_client = TelemetryClientFactory.get_telemetry_client(
		session_id_hex=self.get_session_id_hex()
		)

		driver_connection_params = DriverConnectionParameters(
		http_path=http_path,
		mode=DatabricksClientType.THRIFT,
		host_info=HostDetails(host_url=server_hostname, port=self.port),
		auth_mech=TelemetryHelper.get_auth_mechanism(auth_provider),
		auth_flow=TelemetryHelper.get_auth_flow(auth_provider),
		socket_timeout=kwargs.get("_socket_timeout", None),
		)

		self._telemetry_client.export_initial_telemetry_log(
		driver_connection_params=driver_connection_params,
		user_agent=useragent_header,
		)

		def _set_use_inline_params_with_warning(self, value: Union[bool, str]):
		"""Valid values are True, False, and "silent"

Expand DownExpand Up		@@ -376,7 +419,10 @@ def cursor(
		Will throw an Error if the connection has been closed.
		"""
		if not self.open:
		raise Error("Cannot create cursor from closed connection")
		raise InterfaceError(
		"Cannot create cursor from closed connection",
		session_id_hex=self.get_session_id_hex(),
		)

		cursor = Cursor(
		self,
Expand DownExpand Up		@@ -419,12 +465,17 @@ def _close(self, close_cursors=True) -> None:

		self.open = False

		TelemetryClientFactory.close(self.get_session_id_hex())

		def commit(self):
		"""No-op because Databricks does not support transactions"""
		pass

		def rollback(self):
		raise NotSupportedError("Transactions are not supported on Databricks")
		raise NotSupportedError(
		"Transactions are not supported on Databricks",
		session_id_hex=self.get_session_id_hex(),
		)


		class Cursor:
Expand DownExpand Up		@@ -469,7 +520,10 @@ def __iter__(self):
		for row in self.active_result_set:
		yield row
		else:
		raise Error("There is no active result set")
		raise ProgrammingError(
		"There is no active result set",
		session_id_hex=self.connection.get_session_id_hex(),
		)

		def _determine_parameter_approach(
		self, params: Optional[TParameterCollection]
Expand DownExpand Up		@@ -606,7 +660,10 @@ def _close_and_clear_active_result_set(self):

		def _check_not_closed(self):
		if not self.open:
		raise Error("Attempting operation on closed cursor")
		raise InterfaceError(
		"Attempting operation on closed cursor",
		session_id_hex=self.connection.get_session_id_hex(),
		)

		def _handle_staging_operation(
		self, staging_allowed_local_path: Union[None, str, List[str]]
Expand All		@@ -623,8 +680,9 @@ def _handle_staging_operation(
		elif isinstance(staging_allowed_local_path, type(list())):
		_staging_allowed_local_paths = staging_allowed_local_path
		else:
		raise Error(
		"You must provide at least one staging_allowed_local_path when initialising a connection to perform ingestion commands"
		raise ProgrammingError(
		"You must provide at least one staging_allowed_local_path when initialising a connection to perform ingestion commands",
		session_id_hex=self.connection.get_session_id_hex(),
		)

		abs_staging_allowed_local_paths = [
Expand DownExpand Up		@@ -652,8 +710,9 @@ def _handle_staging_operation(
		else:
		continue
		if not allow_operation:
		raise Error(
		"Local file operations are restricted to paths within the configured staging_allowed_local_path"
		raise ProgrammingError(
		"Local file operations are restricted to paths within the configured staging_allowed_local_path",
		session_id_hex=self.connection.get_session_id_hex(),
		)

		# May be real headers, or could be json string
Expand DownExpand Up		@@ -681,11 +740,13 @@ def _handle_staging_operation(
		handler_args.pop("local_file")
		return self._handle_staging_remove(**handler_args)
		else:
		raiseError(
		raiseProgrammingError(
		f"Operation {row.operation} is not supported. "
		+ "Supported operations are GET, PUT, and REMOVE"
		+ "Supported operations are GET, PUT, and REMOVE",
		session_id_hex=self.connection.get_session_id_hex(),
		)

		@log_latency(StatementType.SQL)
		def _handle_staging_put(
		self, presigned_url: str, local_file: str, headers: Optional[dict] = None
		):
Expand All		@@ -695,7 +756,10 @@ def _handle_staging_put(
		"""

		if local_file is None:
		raise Error("Cannot perform PUT without specifying a local_file")
		raise ProgrammingError(
		"Cannot perform PUT without specifying a local_file",
		session_id_hex=self.connection.get_session_id_hex(),
		)

		with open(local_file, "rb") as fh:
		r = requests.put(url=presigned_url, data=fh, headers=headers)
Expand All		@@ -711,8 +775,9 @@ def _handle_staging_put(
		# fmt: on

		if r.status_code not in [OK, CREATED, NO_CONTENT, ACCEPTED]:
		raise Error(
		f"Staging operation over HTTP was unsuccessful: {r.status_code}-{r.text}"
		raise OperationalError(
		f"Staging operation over HTTP was unsuccessful: {r.status_code}-{r.text}",
		session_id_hex=self.connection.get_session_id_hex(),
		)

		if r.status_code == ACCEPTED:
Expand All		@@ -721,6 +786,7 @@ def _handle_staging_put(
		+ "but not yet applied on the server. It's possible this command may fail later."
		)

		@log_latency(StatementType.SQL)
		def _handle_staging_get(
		self, local_file: str, presigned_url: str, headers: Optional[dict] = None
		):
Expand All		@@ -730,20 +796,25 @@ def _handle_staging_get(
		"""

		if local_file is None:
		raise Error("Cannot perform GET without specifying a local_file")
		raise ProgrammingError(
		"Cannot perform GET without specifying a local_file",
		session_id_hex=self.connection.get_session_id_hex(),
		)

		r = requests.get(url=presigned_url, headers=headers)

		# response.ok verifies the status code is not between 400-600.
		# Any 2xx or 3xx will evaluate r.ok == True
		if not r.ok:
		raise Error(
		f"Staging operation over HTTP was unsuccessful: {r.status_code}-{r.text}"
		raise OperationalError(
		f"Staging operation over HTTP was unsuccessful: {r.status_code}-{r.text}",
		session_id_hex=self.connection.get_session_id_hex(),
		)

		with open(local_file, "wb") as fp:
		fp.write(r.content)

		@log_latency(StatementType.SQL)
		def _handle_staging_remove(
		self, presigned_url: str, headers: Optional[dict] = None
		):
Expand All		@@ -752,10 +823,12 @@ def _handle_staging_remove(
		r = requests.delete(url=presigned_url, headers=headers)

		if not r.ok:
		raise Error(
		f"Staging operation over HTTP was unsuccessful: {r.status_code}-{r.text}"
		raise OperationalError(
		f"Staging operation over HTTP was unsuccessful: {r.status_code}-{r.text}",
		session_id_hex=self.connection.get_session_id_hex(),
		)

		@log_latency(StatementType.QUERY)
		def execute(
		self,
		operation: str,
Expand DownExpand Up		@@ -846,6 +919,7 @@ def execute(

		return self

		@log_latency(StatementType.QUERY)
		def execute_async(
		self,
		operation: str,
Expand DownExpand Up		@@ -951,8 +1025,9 @@ def get_async_execution_result(self):

		return self
		else:
		raise Error(
		f"get_execution_result failed with Operation status {operation_state}"
		raise OperationalError(
		f"get_execution_result failed with Operation status {operation_state}",
		session_id_hex=self.connection.get_session_id_hex(),
		)

		def executemany(self, operation, seq_of_parameters):
Expand All		@@ -970,6 +1045,7 @@ def executemany(self, operation, seq_of_parameters):
		self.execute(operation, parameters)
		return self

		@log_latency(StatementType.METADATA)
		def catalogs(self) -> "Cursor":
		"""
		Get all available catalogs.
Expand All		@@ -993,6 +1069,7 @@ def catalogs(self) -> "Cursor":
		)
		return self

		@log_latency(StatementType.METADATA)
		def schemas(
		self, catalog_name: Optional[str] = None, schema_name: Optional[str] = None
		) -> "Cursor":
Expand DownExpand Up		@@ -1021,6 +1098,7 @@ def schemas(
		)
		return self

		@log_latency(StatementType.METADATA)
		def tables(
		self,
		catalog_name: Optional[str] = None,
Expand DownExpand Up		@@ -1056,6 +1134,7 @@ def tables(
		)
		return self

		@log_latency(StatementType.METADATA)
		def columns(
		self,
		catalog_name: Optional[str] = None,
Expand DownExpand Up		@@ -1102,7 +1181,10 @@ def fetchall(self) -> List[Row]:
		if self.active_result_set:
		return self.active_result_set.fetchall()
		else:
		raise Error("There is no active result set")
		raise ProgrammingError(
		"There is no active result set",
		session_id_hex=self.connection.get_session_id_hex(),
		)

		def fetchone(self) -> Optional[Row]:
		"""
Expand All		@@ -1116,7 +1198,10 @@ def fetchone(self) -> Optional[Row]:
		if self.active_result_set:
		return self.active_result_set.fetchone()
		else:
		raise Error("There is no active result set")
		raise ProgrammingError(
		"There is no active result set",
		session_id_hex=self.connection.get_session_id_hex(),
		)

		def fetchmany(self, size: int) -> List[Row]:
		"""
Expand All		@@ -1138,21 +1223,30 @@ def fetchmany(self, size: int) -> List[Row]:
		if self.active_result_set:
		return self.active_result_set.fetchmany(size)
		else:
		raise Error("There is no active result set")
		raise ProgrammingError(
		"There is no active result set",
		session_id_hex=self.connection.get_session_id_hex(),
		)

		def fetchall_arrow(self) -> "pyarrow.Table":
		self._check_not_closed()
		if self.active_result_set:
		return self.active_result_set.fetchall_arrow()
		else:
		raise Error("There is no active result set")
		raise ProgrammingError(
		"There is no active result set",
		session_id_hex=self.connection.get_session_id_hex(),
		)

		def fetchmany_arrow(self, size) -> "pyarrow.Table":
		self._check_not_closed()
		if self.active_result_set:
		return self.active_result_set.fetchmany_arrow(size)
		else:
		raise Error("There is no active result set")
		raise ProgrammingError(
		"There is no active result set",
		session_id_hex=self.connection.get_session_id_hex(),
		)

		def cancel(self) -> None:
		"""
Expand DownExpand Up		@@ -1455,6 +1549,7 @@ def fetchall_columnar(self):

		return results

		@log_latency()
		def fetchone(self) -> Optional[Row]:
		"""
		Fetch the next row of a query result set, returning a single sequence,
Expand All		@@ -1471,6 +1566,7 @@ def fetchone(self) -> Optional[Row]:
		else:
		return None

		@log_latency()
		def fetchall(self) -> List[Row]:
		"""
		Fetch all (remaining) rows of a query result, returning them as a list of rows.
Expand All		@@ -1480,6 +1576,7 @@ def fetchall(self) -> List[Row]:
		else:
		return self._convert_arrow_table(self.fetchall_arrow())

		@log_latency()
		def fetchmany(self, size: int) -> List[Row]:
		"""
		Fetch the next set of rows of a query result, returning a list of rows.
Expand Down