Jun 3, 2025 · May 25, 2025 · May 25, 2025 · May 25, 2025 · May 25, 2025 · May 25, 2025
diff --git a/src/databricks/sql/backend/databricks_client.py b/src/databricks/sql/backend/databricks_client.py
    from databricks.sql.client import Cursor

 from databricks.sql.thrift_api.TCLIService import ttypes
 from databricks.sql.backend.types import SessionId, CommandId
 from databricks.sql.backend.types import SessionId, CommandId, CommandState
 from databricks.sql.utils import ExecuteResponse
 from databricks.sql.types import SSLOptions

 # Forward reference for type hints
 from typing import TYPE_CHECKING

 if TYPE_CHECKING:
    from databricks.sql.result_set import ResultSet


 class DatabricksClient(ABC):
    # == Connection and Session Management ==
        parameters: List[ttypes.TSparkParameter],
        async_op: bool,
        enforce_embedded_schema_correctness: bool,
    ) ->Optional[ExecuteResponse]:
    ) ->Union["ResultSet", None]:
        """
        Executes a SQL command or query within the specified session.

            enforce_embedded_schema_correctness: Whether to enforce schema correctness

        Returns:
            If async_op is False, returnsan ExecuteResponse object containing the
            If async_op is False, returnsa ResultSet object containing the
            query results and metadata. If async_op is True, returns None and the
            results must be fetched later using get_execution_result().

        pass

    @abstractmethod
    def close_command(self, command_id: CommandId) ->ttypes.TStatus:
    def close_command(self, command_id: CommandId) ->None:
        """
        Closes a command and releases associated resources.

        Args:
            command_id: The command identifier to close

        Returns:
            ttypes.TStatus: The status of the close operation

        Raises:
            ValueError: If the command ID is invalid
            OperationalError: If there's an error closing the command
        """
        pass

    @abstractmethod
    def get_query_state(self, command_id: CommandId) ->ttypes.TOperationState:
    def get_query_state(self, command_id: CommandId) ->CommandState:
        """
        Gets the current state of a query or command.

            command_id: The command identifier to check

        Returns:
 ttypes.TOperationState: The current state of the command
 CommandState: The current state of the command

        Raises:
            ValueError: If the command ID is invalid
        self,
        command_id: CommandId,
        cursor: "Cursor",
    ) ->ExecuteResponse:
    ) ->"ResultSet":
        """
        Retrieves the results of a previously executed command.

            cursor: The cursor object that will handle the results

        Returns:
 ExecuteResponse: An object containing the query results and metadata
 ResultSet: An object containing the query results and metadata

        Raises:
            ValueError: If the command ID is invalid
        max_rows: int,
        max_bytes: int,
        cursor: "Cursor",
    ) ->ExecuteResponse:
    ) ->"ResultSet":
        """
        Retrieves a list of available catalogs.

            cursor: The cursor object that will handle the results

        Returns:
 ExecuteResponse: An object containing the catalog metadata
 ResultSet: An object containing the catalog metadata

        Raises:
            ValueError: If the session ID is invalid
        cursor: "Cursor",
        catalog_name: Optional[str] = None,
        schema_name: Optional[str] = None,
    ) ->ExecuteResponse:
    ) ->"ResultSet":
        """
        Retrieves a list of schemas, optionally filtered by catalog and schema name patterns.

            schema_name: Optional schema name pattern to filter by

        Returns:
 ExecuteResponse: An object containing the schema metadata
 ResultSet: An object containing the schema metadata

        Raises:
            ValueError: If the session ID is invalid
        schema_name: Optional[str] = None,
        table_name: Optional[str] = None,
        table_types: Optional[List[str]] = None,
    ) ->ExecuteResponse:
    ) ->"ResultSet":
        """
        Retrieves a list of tables, optionally filtered by catalog, schema, table name, and table types.

            table_types: Optional list of table types to filter by (e.g., ['TABLE', 'VIEW'])

        Returns:
 ExecuteResponse: An object containing the table metadata
 ResultSet: An object containing the table metadata

        Raises:
            ValueError: If the session ID is invalid
        schema_name: Optional[str] = None,
        table_name: Optional[str] = None,
        column_name: Optional[str] = None,
    ) ->ExecuteResponse:
    ) ->"ResultSet":
        """
        Retrieves a list of columns, optionally filtered by catalog, schema, table, and column name patterns.

            column_name: Optional column name pattern to filter by

        Returns:
 ExecuteResponse: An object containing the column metadata
 ResultSet: An object containing the column metadata

        Raises:
            ValueError: If the session ID is invalid
Original file line number	Diff line number	Diff line change
Expand Up		@@ -15,10 +15,16 @@
		from databricks.sql.client import Cursor

		from databricks.sql.thrift_api.TCLIService import ttypes
		from databricks.sql.backend.types import SessionId, CommandId
		from databricks.sql.backend.types import SessionId, CommandId, CommandState
		from databricks.sql.utils import ExecuteResponse
		from databricks.sql.types import SSLOptions

		# Forward reference for type hints
		from typing import TYPE_CHECKING

		if TYPE_CHECKING:
		from databricks.sql.result_set import ResultSet


		class DatabricksClient(ABC):
		# == Connection and Session Management ==
Expand DownExpand Up		@@ -81,7 +87,7 @@ def execute_command(
		parameters: List[ttypes.TSparkParameter],
		async_op: bool,
		enforce_embedded_schema_correctness: bool,
		) ->Optional[ExecuteResponse]:
		) ->Union["ResultSet", None]:
		"""
		Executes a SQL command or query within the specified session.

Expand All		@@ -101,7 +107,7 @@ def execute_command(
		enforce_embedded_schema_correctness: Whether to enforce schema correctness

		Returns:
		If async_op is False, returnsan ExecuteResponse object containing the
		If async_op is False, returnsa ResultSet object containing the
		query results and metadata. If async_op is True, returns None and the
		results must be fetched later using get_execution_result().

Expand DownExpand Up		@@ -130,7 +136,7 @@ def cancel_command(self, command_id: CommandId) -> None:
		pass

		@abstractmethod
		def close_command(self, command_id: CommandId) ->ttypes.TStatus:
		def close_command(self, command_id: CommandId) ->None:
Copy link Contributor jayantsing-dbJun 1, 2025 Choose a reason for hiding this comment The reason will be displayed to describe this comment to others.Learn more. why changed to void/none? if you choose to return void, make sure the implementations throw appropriate errors for the calling code Copy link ContributorAuthor varun-edachali-dbxJun 2, 2025• edited Loading Choose a reason for hiding this comment The reason will be displayed to describe this comment to others.Learn more. The return type is not used by any of the references, so I chose to convert it to a`None` return type. Any errors in the call will be raised by`make_request`. This structure is identical to the existing`cancel_command` implementation (which returns`None` as well). jayantsing-db reacted with thumbs up emoji
		"""
		Closes a command and releases associated resources.

Expand All		@@ -140,17 +146,14 @@ def close_command(self, command_id: CommandId) -> ttypes.TStatus:
		Args:
		command_id: The command identifier to close

		Returns:
		ttypes.TStatus: The status of the close operation

		Raises:
		ValueError: If the command ID is invalid
		OperationalError: If there's an error closing the command
		"""
		pass

		@abstractmethod
		def get_query_state(self, command_id: CommandId) ->ttypes.TOperationState:
		def get_query_state(self, command_id: CommandId) ->CommandState:
		"""
		Gets the current state of a query or command.

Expand All		@@ -160,7 +163,7 @@ def get_query_state(self, command_id: CommandId) -> ttypes.TOperationState:
		command_id: The command identifier to check

		Returns:
		ttypes.TOperationState: The current state of the command
		CommandState: The current state of the command

		Raises:
		ValueError: If the command ID is invalid
Expand All		@@ -175,7 +178,7 @@ def get_execution_result(
		self,
		command_id: CommandId,
		cursor: "Cursor",
		) ->ExecuteResponse:
		) ->"ResultSet":
varun-edachali-dbx marked this conversation as resolved. Show resolvedHide resolved
		"""
		Retrieves the results of a previously executed command.

Expand All		@@ -187,7 +190,7 @@ def get_execution_result(
		cursor: The cursor object that will handle the results

		Returns:
		ExecuteResponse: An object containing the query results and metadata
		ResultSet: An object containing the query results and metadata

		Raises:
		ValueError: If the command ID is invalid
Expand All		@@ -203,7 +206,7 @@ def get_catalogs(
		max_rows: int,
		max_bytes: int,
		cursor: "Cursor",
		) ->ExecuteResponse:
		) ->"ResultSet":
		"""
		Retrieves a list of available catalogs.

Expand All		@@ -217,7 +220,7 @@ def get_catalogs(
		cursor: The cursor object that will handle the results

		Returns:
		ExecuteResponse: An object containing the catalog metadata
		ResultSet: An object containing the catalog metadata

		Raises:
		ValueError: If the session ID is invalid
Expand All		@@ -234,7 +237,7 @@ def get_schemas(
		cursor: "Cursor",
		catalog_name: Optional[str] = None,
		schema_name: Optional[str] = None,
		) ->ExecuteResponse:
		) ->"ResultSet":
		"""
		Retrieves a list of schemas, optionally filtered by catalog and schema name patterns.

Expand All		@@ -250,7 +253,7 @@ def get_schemas(
		schema_name: Optional schema name pattern to filter by

		Returns:
		ExecuteResponse: An object containing the schema metadata
		ResultSet: An object containing the schema metadata

		Raises:
		ValueError: If the session ID is invalid
Expand All		@@ -269,7 +272,7 @@ def get_tables(
		schema_name: Optional[str] = None,
		table_name: Optional[str] = None,
		table_types: Optional[List[str]] = None,
		) ->ExecuteResponse:
		) ->"ResultSet":
		"""
		Retrieves a list of tables, optionally filtered by catalog, schema, table name, and table types.

Expand All		@@ -287,7 +290,7 @@ def get_tables(
		table_types: Optional list of table types to filter by (e.g., ['TABLE', 'VIEW'])

		Returns:
		ExecuteResponse: An object containing the table metadata
		ResultSet: An object containing the table metadata

		Raises:
		ValueError: If the session ID is invalid
Expand All		@@ -306,7 +309,7 @@ def get_columns(
		schema_name: Optional[str] = None,
		table_name: Optional[str] = None,
		column_name: Optional[str] = None,
		) ->ExecuteResponse:
		) ->"ResultSet":
		"""
		Retrieves a list of columns, optionally filtered by catalog, schema, table, and column name patterns.

Expand All		@@ -324,7 +327,7 @@ def get_columns(
		column_name: Optional column name pattern to filter by

		Returns:
		ExecuteResponse: An object containing the column metadata
		ResultSet: An object containing the column metadata

		Raises:
		ValueError: If the session ID is invalid
Expand Down