add optional schema parameter

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "flit_core.buildapi"
 
 [project]
 name = "vanna"
-version = "0.0.26"
+version = "0.0.27"
 authors = [
   { name="Zain Hoda", email="[email protected]" },
 ]

src/vanna/__init__.py

@@ -1834,6 +1834,7 @@ def connect_to_snowflake(
     username: str,
     password: str,
     database: str,
+    schema: Union[str, None] = None,
     role: Union[str, None] = None,
 ):
     """
@@ -1857,6 +1858,7 @@ def connect_to_snowflake(
         username (str): The Snowflake username.
         password (str): The Snowflake password.
         database (str): The default database to use.
+        schema (Union[str, None], optional): The schema to use. Defaults to None.
         role (Union[str, None], optional): The role to use. Defaults to None.
     """
 
@@ -1914,6 +1916,9 @@ def run_sql_snowflake(sql: str) -> pd.DataFrame:
             cs.execute(f"USE ROLE {role}")
         cs.execute(f"USE DATABASE {database}")
 
+        if schema is not None:
+            cs.execute(f"USE SCHEMA {schema}")
+
         cur = cs.execute(sql)
 
         results = cur.fetchall()