Merge pull request #473 from PucaVaz/main

 feat: Add ClickHouse connection via clickhouse_connect

pyproject.toml

@@ -28,7 +28,7 @@ dependencies = [
 [project.optional-dependencies]
 postgres = ["psycopg2-binary", "db-dtypes"]
 mysql = ["PyMySQL"]
-clickhouse = ["clickhouse_driver"]
+clickhouse = ["clickhouse_connect"]
 bigquery = ["google-cloud-bigquery"]
 snowflake = ["snowflake-connector-python"]
 duckdb = ["duckdb"]

src/vanna/base/base.py

@@ -1023,11 +1023,11 @@ def connect_to_clickhouse(
     ):
 
         try:
-            from clickhouse_driver import connect
+            import clickhouse_connect
         except ImportError:
             raise DependencyError(
                 "You need to install required dependencies to execute this method,"
-                " run command: \npip install clickhouse-driver"
+                " run command: \npip install clickhouse_connect"
             )
 
         if not host:
@@ -1063,32 +1063,30 @@ def connect_to_clickhouse(
         conn = None
 
         try:
-            conn = connect(host=host,
-                                   user=user,
-                                   password=password,
-                                   database=dbname,
-                                   port=port,
-                                  )
+            conn = clickhouse_connect.get_client(
+                host=host,
+                port=port,
+                username=user,
+                password=password,
+                database=dbname,
+            )
             print(conn)
         except Exception as e:
             raise ValidationError(e)
 
         def run_sql_clickhouse(sql: str) -> Union[pd.DataFrame, None]:
             if conn:
                 try:
-                  cs = conn.cursor()
-                  cs.execute(sql)
-                  results = cs.fetchall()
+                    result = conn.query(sql)
+                    results = result.result_rows
 
-                  # Create a pandas dataframe from the results
-                  df = pd.DataFrame(
-                    results, columns=[desc[0] for desc in cs.description]
-                  )
-                  return df
+                    # Create a pandas dataframe from the results
+                    df = pd.DataFrame(results, columns=result.column_names)
+                    return df
 
                 except Exception as e:
                     raise e
-
+        
         self.run_sql_is_set = True
         self.run_sql = run_sql_clickhouse