From 36ddc23193f5247fc293a1b7c8ba87c4759a57f7 Mon Sep 17 00:00:00 2001 From: Clinton Christian Date: Thu, 11 Jul 2024 06:15:50 -0400 Subject: [PATCH] #548: Added additional sqlite args. --- src/vanna/base/base.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/src/vanna/base/base.py b/src/vanna/base/base.py index 4974d064..7c02d216 100644 --- a/src/vanna/base/base.py +++ b/src/vanna/base/base.py @@ -792,13 +792,13 @@ def run_sql_snowflake(sql: str) -> pd.DataFrame: self.run_sql = run_sql_snowflake self.run_sql_is_set = True - def connect_to_sqlite(self, url: str): + def connect_to_sqlite(self, url: str, connect_same_thread: bool = False, **kwargs): """ Connect to a SQLite database. This is just a helper function to set [`vn.run_sql`][vanna.base.base.VannaBase.run_sql] Args: url (str): The URL of the database to connect to. - + connect_same_thread (str): Allow the connection may be accessed in multiple threads. Returns: None """ @@ -817,7 +817,11 @@ def connect_to_sqlite(self, url: str): url = path # Connect to the database - conn = sqlite3.connect(url, check_same_thread=False) + conn = sqlite3.connect( + url, + check_same_thread=connect_same_tread, + **kwargs + ) def run_sql_sqlite(sql: str): return pd.read_sql_query(sql, conn)