Skip to content

Commit

Permalink
fix: arun sql flag
Browse files Browse the repository at this point in the history
  • Loading branch information
vikyw89 committed Jul 5, 2024
1 parent 9345285 commit 0c75730
Showing 1 changed file with 4 additions and 3 deletions.
7 changes: 4 additions & 3 deletions src/vanna/base/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,7 @@ def __init__(self, config=None):

self.config = config
self.run_sql_is_set = False
self.arun_sql_is_set = False
self.static_documentation = ""
self.dialect = self.config.get("dialect", "SQL")
self.language = self.config.get("language", None)
Expand Down Expand Up @@ -1399,7 +1400,7 @@ async def arun_sql_mysql(sql: str, **kwargs) -> pd.DataFrame:
await conn.rollback()
raise e

self.run_sql_is_set = True
self.arun_sql_is_set = True
self.arun_sql = arun_sql_mysql

def connect_to_clickhouse(
Expand Down Expand Up @@ -2003,7 +2004,7 @@ async def arun_sql(self, sql: str, **kwargs) -> pd.DataFrame:
pd.DataFrame: The results of the SQL query.
"""
raise Exception(
"You need to connect to a database first by running vn.connect_to_snowflake(), vn.connect_to_postgres(), similar function, or manually set vn.run_sql"
"You need to connect to a database first by running vn.connect_to_snowflake(), vn.connect_to_postgres(), similar function, or manually set vn.arun_sql"
)

def ask(
Expand Down Expand Up @@ -2172,7 +2173,7 @@ async def aask(
except Exception as e:
print(sql)

if self.run_sql_is_set is False:
if self.arun_sql_is_set is False:
print(
"If you want to run the SQL query, connect to a database first. See here: https://vanna.ai/docs/databases.html"
)
Expand Down

0 comments on commit 0c75730

Please sign in to comment.