add big query

Author: Pinyu Su

application/nlq/data_access/database.py

@@ -15,7 +15,8 @@ class RelationDatabase():
         'starrocks': 'starrocks',
         'clickhouse': 'clickhouse',
         'hive': 'hive',
-        'athena': 'awsathena+rest'
+        'athena': 'awsathena+rest',
+        'bigquery': 'bigquery',
         # Add more mappings here for other databases
     }
 
@@ -37,6 +38,12 @@ def get_db_url(cls, db_type, user, password, host, port, db_name):
                 query={'s3_staging_dir': db_name}
             )
             logger.info(f"db_url: {db_url}")
+        elif db_type == 'bigquery':
+            db_url = db.engine.URL.create(
+                drivername=cls.db_mapping[db_type],
+                host=host,  # BigQuery project. Note: without dataset
+                query={'credentials_path': password}
+            )
         else:
             db_url = db.engine.URL.create(
                 drivername=cls.db_mapping[db_type],
@@ -79,7 +86,7 @@ def get_all_schema_names_by_connection(cls, connection: ConnectConfigEntity):
         if db_type == 'postgresql':
             schemas = [schema for schema in inspector.get_schema_names() if
                        schema not in ('pg_catalog', 'information_schema', 'public')]
-        elif db_type in ('redshift', 'mysql', 'starrocks', 'clickhouse', 'hive', 'athena'):
+        elif db_type in ('redshift', 'mysql', 'starrocks', 'clickhouse', 'hive', 'athena', 'bigquery'):
             schemas = inspector.get_schema_names()
         else:
             raise ValueError("Unsupported database type")
@@ -100,10 +107,14 @@ def get_metadata_by_connection(cls, connection, schemas):
         engine = db.create_engine(db_url)
         # connection = engine.connect()
         metadata = db.MetaData()
-        for s in schemas:
-            metadata.reflect(bind=engine, schema=s, views=True)
-        # metadata.reflect(bind=engine)
-        return metadata
+        if connection.db_type == 'bigquery':
+            metadata.reflect(bind=engine)
+            return metadata
+        else:
+            for s in schemas:
+                metadata.reflect(bind=engine, schema=s, views=True)
+            # metadata.reflect(bind=engine)
+            return metadata
 
     @classmethod
     def get_table_definition_by_connection(cls, connection: ConnectConfigEntity, schemas, table_names):

application/pages/2_🪙_Data_Connection_Management.py

@@ -14,7 +14,8 @@
     'starrocks': 'StarRocks',
     'clickhouse': 'Clickhouse',
     'hive': 'Hive',
-    'athena': 'Athena'
+    'athena': 'Athena',
+    'bigquery': 'BigQuery'
 }
 
 

application/requirements-api.txt

@@ -23,4 +23,5 @@ sqlalchemy-redshift~=0.8.14
 numpy==1.26.4
 pyhive==0.7.0
 thrift==0.20.0
-thrift-sasl==0.4.3
+thrift-sasl==0.4.3
+sqlalchemy-bigquery==1.11.0

application/requirements.txt

@@ -22,4 +22,5 @@ sqlalchemy-redshift~=0.8.14
 numpy==1.26.4
 pyhive==0.7.0
 thrift==0.20.0
-thrift-sasl==0.4.3
+thrift-sasl==0.4.3
+sqlalchemy-bigquery==1.11.0

application/utils/prompt.py

@@ -245,3 +245,12 @@
 The data is：{data}
 
 """
+
+
+BIGQUERY_DIALECT_PROMPT_CLAUDE3 = """
+You are a data analysis expert and proficient in Google BigQuery. Given an input question, first create a syntactically correct BigQuery SQL query to run.
+Unless the user specifies in the question a specific number of examples to obtain, query for at most {top_k} results using the LIMIT clause as per BigQuery.
+Never query for all columns from a table. You must query only the columns that are needed to answer the question. Use backticks (`) to denote table and column names as delimited identifiers.
+Pay attention to use only the column names you can see in the tables below. Be careful to not query for columns that do not exist. Also, pay attention to which column is in which table.
+Pay attention to use CURRENT_DATE() function to get the current date, if the question involves "today". Aside from giving the SQL answer, concisely explain yourself after giving the answer in the same language as the question.
+""".format(top_k=TOP_K)

application/utils/prompts/generate_prompt.py

@@ -1,7 +1,7 @@
 from utils.logging import getLogger
 from utils.prompt import POSTGRES_DIALECT_PROMPT_CLAUDE3, MYSQL_DIALECT_PROMPT_CLAUDE3, \
     DEFAULT_DIALECT_PROMPT, AGENT_COT_EXAMPLE, AWS_REDSHIFT_DIALECT_PROMPT_CLAUDE3, STARROCKS_DIALECT_PROMPT_CLAUDE3, \
-    CLICKHOUSE_DIALECT_PROMPT_CLAUDE3, HIVE_DIALECT_PROMPT_CLAUDE3
+    CLICKHOUSE_DIALECT_PROMPT_CLAUDE3, HIVE_DIALECT_PROMPT_CLAUDE3, BIGQUERY_DIALECT_PROMPT_CLAUDE3
 from utils.prompts import guidance_prompt
 from utils.prompts import table_prompt
 
@@ -2209,6 +2209,8 @@ def generate_llm_prompt(ddl, hints, prompt_map, search_box, sql_examples=None, n
         dialect_prompt = CLICKHOUSE_DIALECT_PROMPT_CLAUDE3
     elif dialect == 'hive':
         dialect_prompt = HIVE_DIALECT_PROMPT_CLAUDE3
+    elif dialect == 'bigquery':
+        dialect_prompt = BIGQUERY_DIALECT_PROMPT_CLAUDE3
     else:
         dialect_prompt = DEFAULT_DIALECT_PROMPT