vanna 0.5.3__py3-none-any.whl → 0.5.4__py3-none-any.whl

vanna/base/base.py

@@ -1012,6 +1012,85 @@ class VannaBase(ABC):
         self.run_sql_is_set = True
         self.run_sql = run_sql_mysql
 
+    def connect_to_clickhouse(
+            self,
+            host: str = None,
+            dbname: str = None,
+            user: str = None,
+            password: str = None,
+            port: int = None,
+    ):
+
+        try:
+            from clickhouse_driver import connect
+        except ImportError:
+            raise DependencyError(
+                "You need to install required dependencies to execute this method,"
+                " run command: \npip install clickhouse-driver"
+            )
+
+        if not host:
+            host = os.getenv("HOST")
+
+        if not host:
+            raise ImproperlyConfigured("Please set your ClickHouse host")
+
+        if not dbname:
+            dbname = os.getenv("DATABASE")
+
+        if not dbname:
+            raise ImproperlyConfigured("Please set your ClickHouse database")
+
+        if not user:
+            user = os.getenv("USER")
+
+        if not user:
+            raise ImproperlyConfigured("Please set your ClickHouse user")
+
+        if not password:
+            password = os.getenv("PASSWORD")
+
+        if not password:
+            raise ImproperlyConfigured("Please set your ClickHouse password")
+
+        if not port:
+            port = os.getenv("PORT")
+
+        if not port:
+            raise ImproperlyConfigured("Please set your ClickHouse port")
+
+        conn = None
+
+        try:
+            conn = connect(host=host,
+                                   user=user,
+                                   password=password,
+                                   database=dbname,
+                                   port=port,
+                                  )
+            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()
+
+                  # Create a pandas dataframe from the results
+                  df = pd.DataFrame(
+                    results, columns=[desc[0] for desc in cs.description]
+                  )
+                  return df
+
+                except Exception as e:
+                    raise e
+
+        self.run_sql_is_set = True
+        self.run_sql = run_sql_clickhouse
+
     def connect_to_oracle(
     self,
     user: str = None,
@@ -1372,6 +1451,10 @@ class VannaBase(ABC):
       def run_sql_presto(sql: str) -> Union[pd.DataFrame, None]:
         if conn:
           try:
+            sql = sql.rstrip()
+            # fix for a known problem with presto db where an extra ; will cause an error.
+            if sql.endswith(';'):
+                sql = sql[:-1]
             cs = conn.cursor()
             cs.execute(sql)
             results = cs.fetchall()
@@ -1393,6 +1476,102 @@ class VannaBase(ABC):
       self.run_sql_is_set = True
       self.run_sql = run_sql_presto
 
+    def connect_to_hive(
+      self,
+      host: str = None,
+      dbname: str = 'default',
+      user: str = None,
+      password: str = None,
+      port: int = None,
+      auth: str = 'CUSTOM'
+    ):
+      """
+        Connect to a Hive database. This is just a helper function to set [`vn.run_sql`][vanna.base.base.VannaBase.run_sql]
+        Connect to a Hive database. This is just a helper function to set [`vn.run_sql`][vanna.base.base.VannaBase.run_sql]
+
+        Args:
+            host (str): The host of the Hive database.
+            dbname (str): The name of the database to connect to.
+            user (str): The username to use for authentication.
+            password (str): The password to use for authentication.
+            port (int): The port to use for the connection.
+            auth (str): The authentication method to use.
+
+        Returns:
+            None
+      """
+
+      try:
+        from pyhive import hive
+      except ImportError:
+        raise DependencyError(
+          "You need to install required dependencies to execute this method,"
+          " run command: \npip install pyhive"
+        )
+
+      if not host:
+        host = os.getenv("HIVE_HOST")
+
+      if not host:
+        raise ImproperlyConfigured("Please set your hive host")
+
+      if not dbname:
+        dbname = os.getenv("HIVE_DATABASE")
+
+      if not dbname:
+        raise ImproperlyConfigured("Please set your hive database")
+
+      if not user:
+        user = os.getenv("HIVE_USER")
+
+      if not user:
+        raise ImproperlyConfigured("Please set your hive user")
+
+      if not password:
+        password = os.getenv("HIVE_PASSWORD")
+
+      if not port:
+        port = os.getenv("HIVE_PORT")
+
+      if not port:
+        raise ImproperlyConfigured("Please set your hive port")
+
+      conn = None
+
+      try:
+        conn = hive.Connection(host=host,
+                               username=user,
+                               password=password,
+                               database=dbname,
+                               port=port,
+                               auth=auth)
+      except hive.Error as e:
+        raise ValidationError(e)
+
+      def run_sql_hive(sql: str) -> Union[pd.DataFrame, None]:
+        if conn:
+          try:
+            cs = conn.cursor()
+            cs.execute(sql)
+            results = cs.fetchall()
+
+            # Create a pandas dataframe from the results
+            df = pd.DataFrame(
+              results, columns=[desc[0] for desc in cs.description]
+            )
+            return df
+
+          except hive.Error as e:
+            print(e)
+            raise ValidationError(e)
+
+          except Exception as e:
+            print(e)
+            raise e
+
+      self.run_sql_is_set = True
+      self.run_sql = run_sql_hive
+
     def run_sql(self, sql: str, **kwargs) -> pd.DataFrame:
         """
         Example:

{vanna-0.5.3.dist-info → vanna-0.5.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vanna
-Version: 0.5.3
+Version: 0.5.4
 Summary: Generate SQL queries from natural language
 Author-email: Zain Hoda <zain@vanna.ai>
 Requires-Python: >=3.9
@@ -41,6 +41,7 @@ Requires-Dist: transformers ; extra == "all"
 Requires-Dist: anthropic ; extra == "anthropic"
 Requires-Dist: google-cloud-bigquery ; extra == "bigquery"
 Requires-Dist: chromadb ; extra == "chromadb"
+Requires-Dist: clickhouse_driver ; extra == "clickhouse"
 Requires-Dist: duckdb ; extra == "duckdb"
 Requires-Dist: google-generativeai ; extra == "gemini"
 Requires-Dist: google-generativeai ; extra == "google"
@@ -68,6 +69,7 @@ Provides-Extra: all
 Provides-Extra: anthropic
 Provides-Extra: bigquery
 Provides-Extra: chromadb
+Provides-Extra: clickhouse
 Provides-Extra: duckdb
 Provides-Extra: gemini
 Provides-Extra: google

{vanna-0.5.3.dist-info → vanna-0.5.4.dist-info}/RECORD

@@ -8,7 +8,7 @@ vanna/ZhipuAI/__init__.py,sha256=NlsijtcZp5Tj9jkOe9fNcOQND_QsGgu7otODsCLBPr0,116
 vanna/anthropic/__init__.py,sha256=85s_2mAyyPxc0T_0JEvYeAkEKWJwkwqoyUwSC5dw9Gk,43
 vanna/anthropic/anthropic_chat.py,sha256=Wk0o-NMW1uvR2fhSWxrR_2FqNh-dLprNG4uuVqpqAkY,2615
 vanna/base/__init__.py,sha256=Sl-HM1RRYzAZoSqmL1CZQmF3ZF-byYTCFQP3JZ2A5MU,28
-vanna/base/base.py,sha256=hUvb94NSVjSIsIa_x38xx8OOdBwL4mmuGmuDj5ednjo,65593
+vanna/base/base.py,sha256=uz0M07nln8nJQfuqq17yCSQOEJF9p1tt8CUGltkcPIM,70850
 vanna/chromadb/__init__.py,sha256=-iL0nW_g4uM8nWKMuWnNePfN4nb9uk8P3WzGvezOqRg,50
 vanna/chromadb/chromadb_vector.py,sha256=eKyPck99Y6Jt-BNWojvxLG-zvAERzLSm-3zY-bKXvaA,8792
 vanna/exceptions/__init__.py,sha256=dJ65xxxZh1lqBeg6nz6Tq_r34jLVmjvBvPO9Q6hFaQ8,685
@@ -37,6 +37,6 @@ vanna/vannadb/__init__.py,sha256=C6UkYocmO6dmzfPKZaWojN0mI5YlZZ9VIbdcquBE58A,48
 vanna/vannadb/vannadb_vector.py,sha256=9YwTO3Lh5owWQE7KPMBqLp2EkiGV0RC1sEYhslzJzgI,6168
 vanna/vllm/__init__.py,sha256=aNlUkF9tbURdeXAJ8ytuaaF1gYwcG3ny1MfNl_cwQYg,23
 vanna/vllm/vllm.py,sha256=QerC3xF5eNzE_nGBDl6YrPYF4WYnjf0hHxxlDWdKX-0,2427
-vanna-0.5.3.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
-vanna-0.5.3.dist-info/METADATA,sha256=JPRFSFUNqsLazznvDMWQIUy_bL9CKRQKCSgPlIwpucU,11248
-vanna-0.5.3.dist-info/RECORD,,
+vanna-0.5.4.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
+vanna-0.5.4.dist-info/METADATA,sha256=kVa5DQIjYZjgyXwVTKQW-2HjGj54ZzAF7cwYUQiXk60,11332
+vanna-0.5.4.dist-info/RECORD,,