Fix: Databricks connection issues

examples/from_databricks.py

@@ -6,22 +6,22 @@
 
 databricks_connector = DatabricksConnector(
     config={
-        "host": "ehxzojy-ue47135",
-        "database": "SNOWFLAKE_SAMPLE_DATA",
-        "token": "",
+        "host": "adb-*****.azuredatabricks.net",
+        "database": "default",
+        "token": "dapidfd412321",
         "port": 443,
-        "table": "lineitem",
-        "httpPath": "tpch_sf1",
+        "table": "loan_payments_data",
+        "httpPath": "/sql/1.0/warehouses/213421312",
         "where": [
             # this is optional and filters the data to
             # reduce the size of the dataframe
-            ["l_quantity", ">", "49"]
+            ["loan_status", "=", "PAIDOFF"],
         ],
     }
 )
 
-llm = OpenAI(api_token="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")
+llm = OpenAI("OPEN_API_KEY")
 df = SmartDataframe(databricks_connector, config={"llm": llm})
 
-response = df.chat("How many records has status 'F'?")
+response = df.chat("How many people from the United states?")
 print(response)

pandasai/connectors/sql.py

@@ -279,10 +279,7 @@ def rows_count(self):
             )
 
         # Run a SQL query to get the number of rows
-        query = sql.text(
-            "SELECT COUNT(*) FROM information_schema.columns "
-            "WHERE table_name = :table_name"
-        ).bindparams(table_name=self._config.table)
+        query = sql.text(f"SELECT COUNT(*) FROM {self._config.table}")
 
         # Return the number of rows
         self._rows_count = self._connection.execute(query).fetchone()[0]
@@ -307,14 +304,7 @@ def columns_count(self):
                 f"{self._config.dialect}"
             )
 
-        # Run a SQL query to get the number of columns
-        query = sql.text(
-            "SELECT COUNT(*) FROM information_schema.columns "
-            f"WHERE table_name = '{self._config.table}'"
-        )
-
-        # Return the number of columns
-        self._columns_count = self._connection.execute(query).fetchone()[0]
+        self._columns_count = len(self.head().columns)
         return self._columns_count
 
     def _get_column_hash(self, include_additional_filters: bool = False):

pyproject.toml

@@ -52,7 +52,7 @@ coverage = "^7.2.7"
 google-cloud-aiplatform = "^1.26.1"
 
 [tool.poetry.extras]
-connectors = ["pymysql", "psycopg2", "snowflake-sqlalchemy", "databricks-sql-connector"]
+connectors = ["pymysql", "psycopg2", "snowflake-sqlalchemy", "sqlalchemy-databricks"]
 google-ai = ["google-generativeai", "google-cloud-aiplatform"]
 google-sheets = ["beautifulsoup4"]
 excel = ["openpyxl"]

tests/connectors/test_databricks.py

@@ -76,11 +76,11 @@ def test_rows_count_property(self):
     def test_columns_count_property(self):
         # Test columns_count property
         self.connector._columns_count = None
-        self.mock_connection.execute.return_value.fetchone.return_value = (
-            8,
-        )  # Sample columns count
+        mock_df = Mock()
+        mock_df.columns = ["Column1", "Column2"]
+        self.connector.head = Mock(return_value=mock_df)
         columns_count = self.connector.columns_count
-        self.assertEqual(columns_count, 8)
+        self.assertEqual(columns_count, 2)
 
     def test_column_hash_property(self):
         # Test column_hash property

tests/connectors/test_snowflake.py

@@ -76,11 +76,11 @@ def test_rows_count_property(self):
     def test_columns_count_property(self):
         # Test columns_count property
         self.connector._columns_count = None
-        self.mock_connection.execute.return_value.fetchone.return_value = (
-            8,
-        )  # Sample columns count
+        mock_df = Mock()
+        mock_df.columns = ["Column1", "Column2"]
+        self.connector.head = Mock(return_value=mock_df)
         columns_count = self.connector.columns_count
-        self.assertEqual(columns_count, 8)
+        self.assertEqual(columns_count, 2)
 
     def test_column_hash_property(self):
         # Test column_hash property

tests/connectors/test_sql.py

@@ -76,11 +76,11 @@ def test_rows_count_property(self):
     def test_columns_count_property(self):
         # Test columns_count property
         self.connector._columns_count = None
-        self.mock_connection.execute.return_value.fetchone.return_value = (
-            8,
-        )  # Sample columns count
+        mock_df = Mock()
+        mock_df.columns = ["Column1", "Column2"]
+        self.connector.head = Mock(return_value=mock_df)
         columns_count = self.connector.columns_count
-        self.assertEqual(columns_count, 8)
+        self.assertEqual(columns_count, 2)
 
     def test_column_hash_property(self):
         # Test column_hash property