Changed the folder structure and tested the basic setup to work

Author: jprakash-db

.idea/databricks-sql-python.iml

@@ -11,7 +11,7 @@
 # # Add the parent directory to sys.path
 # sys.path.append(target_folder_path)
 
-from src.databricks import sql
+from databricks import sql
 
 # from dotenv import load_dotenv
 
@@ -22,7 +22,7 @@
 # load_dotenv()
 
 host = "e2-dogfood.staging.cloud.databricks.com"
-http_path = "/sql/1.0/warehouses/dd43ee29fedd958d"
+http_path = "/sql/1.0/warehouses/58aa1b363649e722"
 
 access_token = ""
 connection = sql.connect(

.idea/vcs.xml

@@ -65,8 +65,8 @@ or to a Databricks Runtime interactive cluster (e.g. /sql/protocolv1/o/123456789
 
 ## Contributing
 
-See [CONTRIBUTING.md](CONTRIBUTING.md)
+See [CONTRIBUTING.md](../CONTRIBUTING.md)
 
 ## License
 
-[Apache License 2.0](LICENSE)
+[Apache License 2.0](../LICENSE)

check.py

@@ -1,7 +1,7 @@
 [tool.poetry]
-name = "databricks-sql-connector"
-version = "3.3.0"
-description = "Databricks SQL Connector for Python"
+name = "databricks-sql-connector-core"
+version = "1.0.0"
+description = "Databricks SQL Connector core for Python"
 authors = ["Databricks <[email protected]>"]
 license = "Apache-2.0"
 readme = "README.md"
@@ -14,15 +14,15 @@ thrift = ">=0.16.0,<0.21.0"
 pandas = [
     { version = ">=1.2.5,<2.2.0", python = ">=3.8" }
 ]
-pyarrow = ">=14.0.1,<17"
+#pyarrow = ">=14.0.1,<17"
 
 lz4 = "^4.0.2"
 requests = "^2.18.1"
 oauthlib = "^3.1.0"
-numpy = [
-    { version = "^1.16.6", python = ">=3.8,<3.11" },
-    { version = "^1.23.4", python = ">=3.11" },
-]
+#numpy = [
+#    { version = "^1.16.6", python = ">=3.8,<3.11" },
+#    { version = "^1.23.4", python = ">=3.11" },
+#]
 sqlalchemy = { version = ">=2.0.21", optional = true }
 openpyxl = "^3.0.10"
 alembic = { version = "^1.0.11", optional = true }
@@ -56,11 +56,11 @@ exclude = ['ttypes\.py$', 'TCLIService\.py$']
 
 [tool.black]
 exclude = '/(\.eggs|\.git|\.hg|\.mypy_cache|\.nox|\.tox|\.venv|\.svn|_build|buck-out|build|dist|thrift_api)/'
-
-[tool.pytest.ini_options]
-markers = {"reviewed" = "Test case has been reviewed by Databricks"}
-minversion = "6.0"
-log_cli = "false"
-log_cli_level = "INFO"
-testpaths = ["tests", "src/databricks/sqlalchemy/test_local"]
-env_files = ["test.env"]
+#
+#[tool.pytest.ini_options]
+#markers = {"reviewed" = "Test case has been reviewed by Databricks"}
+#minversion = "6.0"
+#log_cli = "false"
+#log_cli_level = "INFO"
+#testpaths = ["tests", "src/databricks/sqlalchemy/test_local"]
+#env_files = ["test.env"]

README.md renamed to databricks_sql_connector_core/README.md

@@ -1,7 +1,11 @@
 from typing import Dict, Tuple, List, Optional, Any, Union, Sequence
 
 import pandas
-import pyarrow
+try:
+    import pyarrow
+except ImportError:
+    pyarrow = None
+
 import requests
 import json
 import os
@@ -982,14 +986,14 @@ def fetchmany(self, size: int) -> List[Row]:
         else:
             raise Error("There is no active result set")
 
-    def fetchall_arrow(self) -> pyarrow.Table:
+    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")
 
-    def fetchmany_arrow(self, size) -> pyarrow.Table:
+    def fetchmany_arrow(self, size) -> "pyarrow.Table":
         self._check_not_closed()
         if self.active_result_set:
             return self.active_result_set.fetchmany_arrow(size)
@@ -1160,20 +1164,23 @@ def _convert_arrow_table(self, table):
         # Need to use nullable types, as otherwise type can change when there are missing values.
         # See https://arrow.apache.org/docs/python/pandas.html#nullable-types
         # NOTE: This api is epxerimental https://pandas.pydata.org/pandas-docs/stable/user_guide/integer_na.html
-        dtype_mapping = {
-            pyarrow.int8(): pandas.Int8Dtype(),
-            pyarrow.int16(): pandas.Int16Dtype(),
-            pyarrow.int32(): pandas.Int32Dtype(),
-            pyarrow.int64(): pandas.Int64Dtype(),
-            pyarrow.uint8(): pandas.UInt8Dtype(),
-            pyarrow.uint16(): pandas.UInt16Dtype(),
-            pyarrow.uint32(): pandas.UInt32Dtype(),
-            pyarrow.uint64(): pandas.UInt64Dtype(),
-            pyarrow.bool_(): pandas.BooleanDtype(),
-            pyarrow.float32(): pandas.Float32Dtype(),
-            pyarrow.float64(): pandas.Float64Dtype(),
-            pyarrow.string(): pandas.StringDtype(),
-        }
+        try:
+            dtype_mapping = {
+                pyarrow.int8(): pandas.Int8Dtype(),
+                pyarrow.int16(): pandas.Int16Dtype(),
+                pyarrow.int32(): pandas.Int32Dtype(),
+                pyarrow.int64(): pandas.Int64Dtype(),
+                pyarrow.uint8(): pandas.UInt8Dtype(),
+                pyarrow.uint16(): pandas.UInt16Dtype(),
+                pyarrow.uint32(): pandas.UInt32Dtype(),
+                pyarrow.uint64(): pandas.UInt64Dtype(),
+                pyarrow.bool_(): pandas.BooleanDtype(),
+                pyarrow.float32(): pandas.Float32Dtype(),
+                pyarrow.float64(): pandas.Float64Dtype(),
+                pyarrow.string(): pandas.StringDtype(),
+            }
+        except AttributeError:
+            print("pyarrow is not present")
 
         # Need to rename columns, as the to_pandas function cannot handle duplicate column names
         table_renamed = table.rename_columns([str(c) for c in range(table.num_columns)])
@@ -1190,7 +1197,7 @@ def _convert_arrow_table(self, table):
     def rownumber(self):
         return self._next_row_index
 
-    def fetchmany_arrow(self, size: int) -> pyarrow.Table:
+    def fetchmany_arrow(self, size: int) -> "pyarrow.Table":
         """
         Fetch the next set of rows of a query result, returning a PyArrow table.
 
@@ -1215,7 +1222,7 @@ def fetchmany_arrow(self, size: int) -> pyarrow.Table:
 
         return results
 
-    def fetchall_arrow(self) -> pyarrow.Table:
+    def fetchall_arrow(self) -> "pyarrow.Table":
         """Fetch all (remaining) rows of a query result, returning them as a PyArrow table."""
         results = self.results.remaining_rows()
         self._next_row_index += results.num_rows

pyproject.toml renamed to databricks_sql_connector_core/pyproject.toml

@@ -8,7 +8,10 @@
 from ssl import CERT_NONE, CERT_REQUIRED, create_default_context
 from typing import List, Union
 
-import pyarrow
+try:
+    import pyarrow
+except ImportError:
+    pyarrow = None
 import thrift.transport.THttpClient
 import thrift.protocol.TBinaryProtocol
 import thrift.transport.TSocket
@@ -37,7 +40,7 @@
     convert_column_based_set_to_arrow_table,
 )
 
-from src.databricks.sql.thrift_api.TCLIService.ttypes import TDBSqlResultFormat
+# from databricks.sql import TDBSqlResultFormat
 
 logger = logging.getLogger(__name__)
 
@@ -654,6 +657,10 @@ def _get_metadata_resp(self, op_handle):
 
     @staticmethod
     def _hive_schema_to_arrow_schema(t_table_schema):
+
+        if pyarrow is None:
+            raise ImportError("pyarrow is required to convert Hive schema to Arrow schema")
+
         def map_type(t_type_entry):
             if t_type_entry.primitiveEntry:
                 return {
@@ -760,12 +767,17 @@ def _results_message_to_execute_response(self, resp, operation_state):
         description = self._hive_schema_to_description(
             t_result_set_metadata_resp.schema
         )
-        schema_bytes = (
-            t_result_set_metadata_resp.arrowSchema
-            or self._hive_schema_to_arrow_schema(t_result_set_metadata_resp.schema)
-            .serialize()
-            .to_pybytes()
-        )
+
+        if pyarrow:
+            schema_bytes = (
+                t_result_set_metadata_resp.arrowSchema
+                or self._hive_schema_to_arrow_schema(t_result_set_metadata_resp.schema)
+                .serialize()
+                .to_pybytes()
+            )
+        else:
+            schema_bytes = None
+
         lz4_compressed = t_result_set_metadata_resp.lz4Compressed
         is_staging_operation = t_result_set_metadata_resp.isStagingOperation
         if direct_results and direct_results.resultSet:

src/databricks/__init__.py renamed to databricks_sql_connector_core/src/databricks/__init__.py

@@ -14,7 +14,10 @@
 from ssl import SSLContext
 
 import lz4.frame
-import pyarrow
+try:
+    import pyarrow
+except ImportError:
+    pyarrow = None
 
 from databricks.sql import OperationalError, exc
 from databricks.sql.cloudfetch.download_manager import ResultFileDownloadManager
@@ -155,7 +158,7 @@ def remaining_rows(self):
 class ArrowQueue(ResultSetQueue):
     def __init__(
         self,
-        arrow_table: pyarrow.Table,
+        arrow_table: "pyarrow.Table",
         n_valid_rows: int,
         start_row_index: int = 0,
     ):
@@ -170,7 +173,7 @@ def __init__(
         self.arrow_table = arrow_table
         self.n_valid_rows = n_valid_rows
 
-    def next_n_rows(self, num_rows: int) -> pyarrow.Table:
+    def next_n_rows(self, num_rows: int) -> 'pyarrow.Table':
         """Get upto the next n rows of the Arrow dataframe"""
         length = min(num_rows, self.n_valid_rows - self.cur_row_index)
         # Note that the table.slice API is not the same as Python's slice
@@ -179,7 +182,7 @@ def next_n_rows(self, num_rows: int) -> pyarrow.Table:
         self.cur_row_index += slice.num_rows
         return slice
 
-    def remaining_rows(self) -> pyarrow.Table:
+    def remaining_rows(self) -> 'pyarrow.Table':
         slice = self.arrow_table.slice(
             self.cur_row_index, self.n_valid_rows - self.cur_row_index
         )
@@ -239,7 +242,7 @@ def __init__(
         self.table = self._create_next_table()
         self.table_row_index = 0
 
-    def next_n_rows(self, num_rows: int) -> pyarrow.Table:
+    def next_n_rows(self, num_rows: int) -> 'pyarrow.Table':
         """
         Get up to the next n rows of the cloud fetch Arrow dataframes.
 
@@ -271,7 +274,7 @@ def next_n_rows(self, num_rows: int) -> pyarrow.Table:
         logger.debug("CloudFetchQueue: collected {} next rows".format(results.num_rows))
         return results
 
-    def remaining_rows(self) -> pyarrow.Table:
+    def remaining_rows(self) -> 'pyarrow.Table':
         """
         Get all remaining rows of the cloud fetch Arrow dataframes.
 
@@ -292,7 +295,7 @@ def remaining_rows(self) -> pyarrow.Table:
             self.table_row_index = 0
         return results
 
-    def _create_next_table(self) -> Union[pyarrow.Table, None]:
+    def _create_next_table(self) -> Union['pyarrow.Table', None]:
         logger.debug(
             "CloudFetchQueue: Trying to get downloaded file for row {}".format(
                 self.start_row_index
@@ -331,7 +334,7 @@ def _create_next_table(self) -> Union[pyarrow.Table, None]:
 
         return arrow_table
 
-    def _create_empty_table(self) -> pyarrow.Table:
+    def _create_empty_table(self) -> 'pyarrow.Table':
         # Create a 0-row table with just the schema bytes
         return create_arrow_table_from_arrow_file(self.schema_bytes, self.description)
 
@@ -570,7 +573,7 @@ def transform_paramstyle(
     return output
 
 
-def create_arrow_table_from_arrow_file(file_bytes: bytes, description) -> pyarrow.Table:
+def create_arrow_table_from_arrow_file(file_bytes: bytes, description) -> 'pyarrow.Table':
     arrow_table = convert_arrow_based_file_to_arrow_table(file_bytes)
     return convert_decimals_in_arrow_table(arrow_table, description)
 
@@ -597,7 +600,7 @@ def convert_arrow_based_set_to_arrow_table(arrow_batches, lz4_compressed, schema
     return arrow_table, n_rows
 
 
-def convert_decimals_in_arrow_table(table, description) -> pyarrow.Table:
+def convert_decimals_in_arrow_table(table, description) -> 'pyarrow.Table':
     for i, col in enumerate(table.itercolumns()):
         if description[i][1] == "decimal":
             decimal_col = col.to_pandas().apply(

src/databricks/sql/__init__.py renamed to databricks_sql_connector_core/src/databricks/sql/__init__.py

@@ -0,0 +1,5 @@
+try:
+    from databricks_sqlalchemy import *
+except:
+    import warnings
+    warnings.warn("Install databricks-sqlalchemy plugin before using this")