I have tried the solutions. But they don't work.

When i tried :

from pyspark.sql import SparkSession
from databricks import sql

def get_db(_=None):
db = SQLDatabase.from_databricks(...)

spark = SparkSession.builder.appName("NewSession").getOrCreate()
sql_context = sql.Context(spark.sparkContext)
sql_context.sessionState().newSession()

return db

It returned : AttributeError: module 'databricks.sql' has no attribute 'Context'.

When I tried :

def get_chain():
global db
if db is None or not db.is_valid():
db = get_db()
chain = SQLDatabaseChain.from_llm(llm=llm, db=db, verbose=True)
return chain

It returned: AttributeError: 'SQLDatabase' object has no attribute 'is_valid'.