Refractored the test code and moved to respective folders

Author: jprakash-db

databricks_sql_connector_core/src/databricks_sql_connector_core/sql/client.py

@@ -784,8 +784,6 @@ def execute(
             parameters=prepared_params,
         )
 
-        # print("Line 781")
-        # print(execute_response)
         self.active_result_set = ResultSet(
             self.connection,
             execute_response,
@@ -1141,7 +1139,6 @@ def _fill_results_buffer(self):
     def _convert_columnar_table(self, table):
         column_names = [c[0] for c in self.description]
         ResultRow = Row(*column_names)
-        # print("Table\n",table)
         result = []
         for row_index in range(len(table[0])):
             curr_row = []

databricks_sql_connector_core/src/databricks_sql_connector_core/sql/thrift_backend.py

@@ -743,7 +743,6 @@ def _results_message_to_execute_response(self, resp, operation_state):
         else:
             t_result_set_metadata_resp = self._get_metadata_resp(resp.operationHandle)
 
-        # print(f"Line 739 - {t_result_set_metadata_resp.resultFormat}")
         if t_result_set_metadata_resp.resultFormat not in [
             ttypes.TSparkRowSetType.ARROW_BASED_SET,
             ttypes.TSparkRowSetType.COLUMN_BASED_SET,
@@ -880,18 +879,8 @@ def execute_command(
                 # We want to receive proper Timestamp arrow types.
                 "spark.thriftserver.arrowBasedRowSet.timestampAsString": "false"
             },
-            # useArrowNativeTypes=spark_arrow_types,
-            # canReadArrowResult=True,
-            # # canDecompressLZ4Result=lz4_compression,
-            # canDecompressLZ4Result=False,
-            # canDownloadResult=False,
-            # # confOverlay={
-            # #     # We want to receive proper Timestamp arrow types.
-            # #     "spark.thriftserver.arrowBasedRowSet.timestampAsString": "false"
-            # # },
-            # resultDataFormat=TDBSqlResultFormat(None,None,True),
-            # # useArrowNativeTypes=spark_arrow_types,
-             parameters=parameters,
+            useArrowNativeTypes=spark_arrow_types,
+            parameters=parameters,
         )
         resp = self.make_request(self._client.ExecuteStatement, req)
         return self._handle_execute_response(resp, cursor)

databricks_sql_connector_core/src/databricks_sql_connector_core/sql/utils.py

@@ -74,18 +74,6 @@ def build_queue(
         Returns:
             ResultSetQueue
         """
-
-        # def trow_to_json(trow):
-        #     # Step 1: Serialize TRow using Thrift's TJSONProtocol
-        #     transport = TTransport.TMemoryBuffer()
-        #     protocol = TJSONProtocol.TJSONProtocol(transport)
-        #     trow.write(protocol)
-        #
-        #     # Step 2: Extract JSON string from the transport
-        #     json_str = transport.getvalue().decode('utf-8')
-        #
-        #     return json_str
-
         if row_set_type == TSparkRowSetType.ARROW_BASED_SET:
             arrow_table, n_valid_rows = convert_arrow_based_set_to_arrow_table(
                 t_row_set.arrowBatches, lz4_compressed, arrow_schema_bytes
@@ -95,30 +83,11 @@ def build_queue(
             )
             return ArrowQueue(converted_arrow_table, n_valid_rows)
         elif row_set_type == TSparkRowSetType.COLUMN_BASED_SET:
-            # print("Lin 79 ")
-            # print(type(t_row_set))
-            # print(t_row_set)
-            # json_str = json.loads(trow_to_json(t_row_set))
-            # pretty_json = json.dumps(json_str, indent=2)
-            # print(pretty_json)
-
             converted_column_table, column_names = convert_column_based_set_to_column_table(
                 t_row_set.columns,
                 description)
-            # print(converted_column_table, column_names)
 
             return ColumnQueue(converted_column_table, column_names)
-
-            # print(columnQueue.next_n_rows(2))
-            # print(columnQueue.next_n_rows(2))
-            # print(columnQueue.remaining_rows())
-            # arrow_table, n_valid_rows = convert_column_based_set_to_arrow_table(
-            #     t_row_set.columns, description
-            # )
-            # converted_arrow_table = convert_decimals_in_arrow_table(
-            #     arrow_table, description
-            # )
-            # return ArrowQueue(converted_arrow_table, n_valid_rows)
         elif row_set_type == TSparkRowSetType.URL_BASED_SET:
             return CloudFetchQueue(
                 schema_bytes=arrow_schema_bytes,

tests/__init__.py renamed to databricks_sql_connector_core/tests/__init__.py

@@ -8,13 +8,13 @@
 
 
 def pysql_supports_arrow():
-    """Import databricks.sql and test whether Cursor has fetchall_arrow."""
-    from databricks.sql import Cursor
+    """Import databricks_sql_connector_core.sql and test whether Cursor has fetchall_arrow."""
+    from databricks_sql_connector_core.sql.client import Cursor
     return hasattr(Cursor, 'fetchall_arrow')
 
 
 def pysql_has_version(compare, version):
-    """Import databricks.sql, and return compare_module_version(...).
+    """Import databricks_sql_connector_core.sql, and return compare_module_version(...).
 
     Expected use:
         from common.predicates import pysql_has_version
@@ -98,4 +98,4 @@ def validate_version(version):
 
     mod_version = validate_version(module.__version__)
     req_version = validate_version(version)
-    return compare_versions(compare, mod_version, req_version)
+    return compare_versions(compare, mod_version, req_version)