How to cache a function that takes a database connection as an input using streamlit

Question

I have a database connection:

import sqlalchemy as sa

engine = sa.create_engine('my_info')
connection = engine.connect()

Subsequently I have a function:

import pandas as pd

def load_data(connection):
    sql = 'select * from tablename'
    df = pd.read_sql(sql, con = connection)
    return df

This is part of an app in Streamlit that I'm working on, I need streamlit to cache the output of my load_data function. I thought it worked like this:

import pandas as pd
import streamlit as st

@st.cache()
def load_data(connection):
    sql = 'select * from tablename'
    df = pd.read_sql(sql, con = connection)
    return df

But this gives me the following error:

UnhashableTypeError: Cannot hash object of type builtins.weakref, found in the arguments of load_data().

The error is much longer, and if it helps I will post it. The error also contains a link to the streamlit documentation. I read it and reformulated my code to look like this:

@st.cache()
def DBConnection():
    engine = sa.create_engine("my_info")
    conn = engine.connect()
    return conn
conn = DBConnection()

@st.cache(hash_funcs={DBConnection: id})
def load_data(connection):
        sql = 'select * from tablename'
        df = pd.read_sql(sql, con = connection)
        return df

But this gives me a NameError:

NameError: name 'DBConnection' is not defined

I've run out of idea's to try, any help would be highly appreciated. It is very possible that I misunderstood the documentation as it assumes a lot of prior knowledge about the process of hashing and caching.

Answer 1

Combine the two methods and use:

@st.cache(allow_output_mutation=true)

Code:

@st.cache(allow_output_mutation=true)
def load_data():
    engine = sa.create_engine("my_info")
    conn = engine.connect()
    sql = 'select * from tablename'
    df = pd.read_sql(sql, con = conn )
    return df

For more you can read in documentation