vanna-ai / vanna

🤖 Chat with your SQL database 📊. Accurate Text-to-SQL Generation via LLMs using RAG 🔄.
https://vanna.ai/docs/
MIT License
12.13k stars 972 forks source link

BigQuery Connection Error 🐛 #348

Open ProdigyMaster opened 7 months ago

ProdigyMaster commented 7 months ago

Describe the bug Hey, I'm having an issue when Vanna tries to run a query in BigQuery. The error message says: "('invalid_grant: Bad Request', {'error': 'invalid_grant', 'error_description': 'Bad Request'})".

I use the same credentials json keyfile, and it works well with my other python scripts.

image

To Reproduce Steps to reproduce the behavior:

  1. Try to connect to BigQuery using json keyfile: vn.connect_to_bigquery( project_id= "XYZ", cred_file_path= r'XYZ/keyfile.json', )

  2. Go to VannaFlaskApp(vn)

  3. See error when query

Expected behavior I expected this method to work with the keyfile.json

Error logs/Screenshots image

Same with console: Couldn't run sql: ('invalid_grant: Bad Request', {'error': 'invalid_grant', 'error_description': 'Bad Request'})

Desktop (please complete the following information where):

zainhoda commented 7 months ago

Are you able to use that same keyfile to run your BigQuery query outside of Vanna? One possibility is that the credentials allow other types of access but not specifically BigQuery queries (or perhaps queries against the specific dataset).

I think the easiest way to both verify this and to get you a quick solution regardless is for you to create a Python function called run_sql_bigquery that takes in sql and returns a dataframe using your credential file, i.e.:

def run_sql_bigquery(sql: str) -> pd.DataFrame

If you're able to get that the successfully run, then you can actually just do this to give Vanna access to that function.

vn.run_sql = run_sql_bigquery
vn.run_sql_is_set = True
ProdigyMaster commented 7 months ago

@zainhoda You're awesome! :) Yes, it works by using the run function and native libraries to connect to BigQuery. So, the issue seems to be with the function in Vanna:

vn.connect_to_bigquery(
    project_id="XYZ",
    cred_file_path=r'XYZ/keyfile.json'
)
ProdigyMaster commented 7 months ago

Here is the code that works and might be helpful as a reference:

import os
from google.cloud import bigquery
from vanna.openai.openai_chat import OpenAI_Chat
from vanna.chromadb.chromadb_vector import ChromaDB_VectorStore
from vanna.flask import VannaFlaskApp

class MyVanna(ChromaDB_VectorStore, OpenAI_Chat):
    def __init__(self, config=None):
        ChromaDB_VectorStore.__init__(self, config=config)
        OpenAI_Chat.__init__(self, config=config)

vn = MyVanna(config={})

service_account_path = r'...\keyfile.json'
os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = service_account_path

client = bigquery.Client()

def run_sql(sql: str) -> pd.DataFrame:
    query_job = client.query(sql)
    results = query_job.result()
    df = results.to_dataframe()
    return df

vn.run_sql = run_sql
vn.run_sql_is_set = True