Queries don't properly execute without "REBOOT" in Vanna Streamlit

[sathyapriyaa13]

Description of the bug When we prompt a question to vanna, we get the query and when it runs, it produces psycopg2 interface error. While the same query after REBOOT produces expected results.

Error logs/Screenshots psycopg2.InterfaceError: connection already closed

2024-02-26 07:28:46.296 Uncaught app exception

Traceback (most recent call last):

File "/home/adminuser/venv/lib/python3.9/site-packages/vanna/base/base.py", line 721, in run_sql_postgres

cs = conn.cursor()

psycopg2.InterfaceError: connection already closed

During handling of the above exception, another exception occurred:

Traceback (most recent call last):

File "/home/adminuser/venv/lib/python3.9/site-packages/streamlit/runtime/scriptrunner/script_runner.py", line 535, in _run_script

exec(code, module.__dict__)

File "/mount/src/jules_chat_demo/app.py", line 72, in

df = run_sql_cached(sql=sql)

File "/home/adminuser/venv/lib/python3.9/site-packages/streamlit/runtime/caching/cache_utils.py", line 212, in wrapper

return cached_func(*args, **kwargs)

File "/home/adminuser/venv/lib/python3.9/site-packages/streamlit/runtime/caching/cache_utils.py", line 241, in call

return self._get_or_create_cached_value(args, kwargs)

File "/home/adminuser/venv/lib/python3.9/site-packages/streamlit/runtime/caching/cache_utils.py", line 268, in _get_or_create_cached_value

return self._handle_cache_miss(cache, value_key, func_args, func_kwargs)

File "/home/adminuser/venv/lib/python3.9/site-packages/streamlit/runtime/caching/cache_utils.py", line 324, in _handle_cache_miss

computed_value = self._info.func(*func_args, **func_kwargs)

File "/mount/src/jules_chat_demo/utils/vanna_calls.py", line 23, in run_sql_cached

return vn.run_sql(sql=sql)

File "/home/adminuser/venv/lib/python3.9/site-packages/vanna/base/base.py", line 732, in run_sql_postgres

conn.rollback()

psycopg2.InterfaceError: connection already closed @zainhoda -- This is the issue we face in vanna [streamlit], the same query runs fine after "REBOOT" but before "REBOOT" this is the condition, could you please check?

Expected behavior The answers are expected without developers "REBOOTING" on facing this issue as this leads to bad ux.

Desktop (please complete the following information where):

• OS: [e.g. Windows]
• Version: [e.g. 10]
• Python: [3.9]
• Vanna(Streamlit): [0.1.1]

[akshatdalton]

@zainhoda @andreped any update on this? We allow our developers to switch to any db to execute the query (user prompt input) but recently I received the error:

got error: connection already closed

when the connection was established with PostgreSQL DB.

I had this connection set up in a pod that was ~7days old.

I think postgres or a similar db server closes the connection if it is idle for too long.

Vanna needs to handle such cases and try to re-establish the connection accordingly.

[zainhoda]

@akshatdalton as a first pass, instead of doing vn.connect_to_postgres could you try something like this:

import psycopg2
import pandas as pd
from psycopg2 import Error
from typing import Union

def run_sql_postgres(sql: str) -> Union[pd.DataFrame, None]:
    try:
        # Establish connection
        with psycopg2.connect(
                host=host,
                dbname=dbname,
                user=user,
                password=password,
                port=port,
            ) as conn:

            # Execute the SQL statement
            with conn.cursor() as cs:
                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 Error as e:
        print(f"Database error: {e}")
        raise ValidationError(e)

    except Exception as e:
        print(f"Unexpected error: {e}")
        raise e

vn.dialect = "PostgreSQL"
vn.run_sql_is_set = True
vn.run_sql = run_sql_postgres

Another option to try is:

from psycopg2 import pool

pool = pool.Pooling(
    minconn=1, maxconn=5, # Adjust min/max connections as needed
    host=host,
    dbname=dbname,
    user=user,
    password=password,
    port=port,
)

def run_sql_postgres(sql: str) -> Union[pd.DataFrame, None]:
    with pool.getconn() as conn:
        cs = conn.cursor()
        # ... rest of your code using cursor

Please report back what works for you and we can incorporate into the package.

[akshatdalton]

Hi @zainhoda, I believe both of the codes work; but the above code adds additional latency/resource consumption every time we execute a sql query.

Instead, we should catch psycopg2.OperationalError (catch similar exceptions for other dbs) if the connection is closed and try to recreate the conn object.

The idea is to create conn obj once and create a cusor everytime we execute an SQL query.

Let me know your thoughts.

[akshatdalton]

In my org, we support multiple dbs for our customers and hence it becomes very necessary how we create and handle the db connections.

[akshatdalton]

Hi @zainhoda, any update on this?

[zainhoda]

@akshatdalton were you able to create a customized run_sql function that works for you?

[sathyapriyaa13]

any updates on this?

On Mon, May 20, 2024 at 11:25 PM Zain Hoda @.***> wrote:

@akshatdalton https://github.com/akshatdalton were you able to create a customized run_sql function that works for you?

— Reply to this email directly, view it on GitHub https://github.com/vanna-ai/vanna/issues/261#issuecomment-2120927836, or unsubscribe https://github.com/notifications/unsubscribe-auth/BDEQMOFD4YVQTD4EIGYJI7TZDI2HTAVCNFSM6AAAAABDZWAWCWVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDCMRQHEZDOOBTGY . You are receiving this because you authored the thread.Message ID: @.***>

[hotsun]

I can see the latest vanna package added **kwargs for psycopg2.connect(), you can try to add the keealive parameters for vn.connect_to_postgres(), eg: vn.connect_to_postgres( host=db_config['host'], dbname=db_config['dbname'], user=db_config['username'], password=db_config['password'], port=db_config['port'], keepalives=1, keepalives_idle=30, keepalives_interval=180 )

Ref: https://www.psycopg.org/docs/module.html https://www.postgresql.org/docs/current/libpq-connect.html#LIBPQ-KEEPALIVES

[sathyapriyaa13]

Thanks!

On Thu, 15 Aug, 2024, 11:06 am Sunny Kwok, @.***> wrote:

I can see the latest vanna package added **kwargs for psycopg2.connect(), you can try to add the keealive parameters for vn.connect_to_postgres(), eg: vn.connect_to_postgres( host=db_config['host'], dbname=db_config['dbname'], user=db_config['username'], password=db_config['password'], port=db_config['port'], keepalives=1, keepalives_idle=30, keepalives_interval=180, )

Ref: https://www.psycopg.org/docs/module.html https://www.postgresql.org/docs/current/libpq-connect.html#LIBPQ-KEEPALIVES

— Reply to this email directly, view it on GitHub https://github.com/vanna-ai/vanna/issues/261#issuecomment-2290678107, or unsubscribe https://github.com/notifications/unsubscribe-auth/BDEQMOEIVLHNY5CYH3UVTMDZRQ455AVCNFSM6AAAAABDZWAWCWVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDEOJQGY3TQMJQG4 . You are receiving this because you authored the thread.Message ID: @.***>