Merge pull request #314 from aws-samples/v1.8.0_dev

add bigquery support

Author: supinyu

application/nlq/business/connection.py

@@ -53,6 +53,11 @@ def get_db_url_by_name(cls, conn_name):
         conn_config = cls.get_conn_config_by_name(conn_name)
         return RelationDatabase.get_db_url_by_connection(conn_config)
 
+    @classmethod
+    def get_db_password_host_by_name(cls, conn_name):
+        conn_config = cls.get_conn_config_by_name(conn_name)
+        return RelationDatabase.get_password_host_by_connection(conn_config)
+
     @classmethod
     def get_db_type_by_name(cls, conn_name):
         conn_config = cls.get_conn_config_by_name(conn_name)

application/nlq/data_access/database.py

@@ -1,3 +1,4 @@
+import json
 
 import sqlalchemy as db
 from sqlalchemy import text, Column, inspect
@@ -39,10 +40,11 @@ def get_db_url(cls, db_type, user, password, host, port, db_name):
             )
             logger.info(f"db_url: {db_url}")
         elif db_type == 'bigquery':
+            password = json.loads(password)
             db_url = db.engine.URL.create(
                 drivername=cls.db_mapping[db_type],
                 host=host,  # BigQuery project. Note: without dataset
-                query={'credentials_path': password}
+                query={'credentials_path': json.dumps(password)}
             )
         else:
             db_url = db.engine.URL.create(
@@ -67,7 +69,11 @@ def test_connection(cls, db_type, user, password, host, port, db_name) -> bool:
                     connect_args={'s3_staging_dir': db_name}
                 )
             else:
-                engine = db.create_engine(cls.get_db_url(db_type, user, password, host, port, db_name))
+                if db_type == "bigquery":
+                    password = json.loads(password)
+                    engine = db.create_engine(url=host, credentials_info=password)
+                else:
+                    engine = db.create_engine(cls.get_db_url(db_type, user, password, host, port, db_name))
             connection = engine.connect()
             return True
         except Exception as e:
@@ -80,7 +86,11 @@ def get_all_schema_names_by_connection(cls, connection: ConnectConfigEntity):
         db_type = connection.db_type
         db_url = cls.get_db_url(db_type, connection.db_user, connection.db_pwd, connection.db_host, connection.db_port,
                                 connection.db_name)
-        engine = db.create_engine(db_url)
+        if db_type == "bigquery":
+            password = json.loads(connection.db_pwd)
+            engine = db.create_engine(url=connection.db_host, credentials_info=password)
+        else:
+            engine = db.create_engine(db_url)
         inspector = inspect(engine)
 
         if db_type == 'postgresql':
@@ -104,7 +114,11 @@ def get_all_tables_by_connection(cls, connection: ConnectConfigEntity, schemas=N
     def get_metadata_by_connection(cls, connection, schemas):
         db_url = cls.get_db_url(connection.db_type, connection.db_user, connection.db_pwd, connection.db_host,
                                 connection.db_port, connection.db_name)
-        engine = db.create_engine(db_url)
+        if connection.db_type == "bigquery":
+            password = json.loads(connection.db_pwd)
+            engine = db.create_engine(url=connection.db_host, credentials_info=password)
+        else:
+            engine = db.create_engine(db_url)
         # connection = engine.connect()
         metadata = db.MetaData()
         if connection.db_type == 'bigquery':
@@ -179,3 +193,7 @@ def get_db_url_by_connection(cls, connection: ConnectConfigEntity):
         db_url = cls.get_db_url(connection.db_type, connection.db_user, connection.db_pwd, connection.db_host,
                                 connection.db_port, connection.db_name)
         return db_url
+
+    @classmethod
+    def get_password_host_by_connection(cls, connection: ConnectConfigEntity):
+        return connection.db_pwd, connection.db_host

application/pages/2_🪙_Data_Connection_Management.py

@@ -1,10 +1,13 @@
+import json
+
 import streamlit as st
 from dotenv import load_dotenv
 from nlq.business.connection import ConnectionManagement
 from nlq.data_access.database import RelationDatabase
+from utils.logging import getLogger
 from utils.navigation import make_sidebar
 
-
+logger = getLogger()
 # global variables
 
 db_type_mapping = {
@@ -78,23 +81,39 @@ def main():
         db_type = db_type.lower()  # Convert to lowercase for matching with db_mapping keys
         if db_type == 'athena':
             st.info("Please enter S3 staging directory in the database name field. You can leave other fields empty. Please also make sure that IAM role is able to access Athena and S3.")
-        host = st.text_input("Enter host")
-        port = st.text_input("Enter port")
-        user = st.text_input("Enter username")
-        password = st.text_input("Enter password", type="password")
-        db_name = st.text_input("Enter database name")
-        comment = st.text_input("Enter comment")
-
-        test_connection_view(db_type, user, password, host, port, db_name)
 
-        if st.button('Add Connection', type='primary'):
-            if db_name == '':
-                st.error("Database name is required!")
-            else:
+        if db_type == "bigquery":
+            host = st.text_input("Enter host")
+            password = st.text_area("Credentials Info", height=200,  placeholder="Paste your credentials info here")
+            # credentials_info = json.loads(credentials_info)
+            port = ""
+            user = ""
+            db_name = ""
+            comment = st.text_input("Enter comment")
+            test_connection_view(db_type, user, password, host, port, db_name)
+            if st.button('Add Connection', type='primary'):
                 ConnectionManagement.add_connection(connection_name, db_type, host, port, user, password, db_name, comment)
                 st.success(f"{connection_name} added successfully!")
                 st.session_state.new_connection_mode = False
 
+        else:
+            host = st.text_input("Enter host")
+            port = st.text_input("Enter port")
+            user = st.text_input("Enter username")
+            password = st.text_input("Enter password", type="password")
+            db_name = st.text_input("Enter database name")
+            comment = st.text_input("Enter comment")
+
+            test_connection_view(db_type, user, password, host, port, db_name)
+
+            if st.button('Add Connection', type='primary'):
+                if db_name == '':
+                    st.error("Database name is required!")
+                else:
+                    ConnectionManagement.add_connection(connection_name, db_type, host, port, user, password, db_name, comment)
+                    st.success(f"{connection_name} added successfully!")
+                    st.session_state.new_connection_mode = False
+
     elif st.session_state.update_connection_mode:
         st.subheader("Update Database Connection")
         current_conn = st.session_state.current_connection

application/utils/apis.py

@@ -1,3 +1,5 @@
+import json
+
 import sqlalchemy as db
 from sqlalchemy import text
 from utils.env_var import RDS_MYSQL_HOST, RDS_MYSQL_PORT, RDS_MYSQL_USERNAME, RDS_MYSQL_PASSWORD, RDS_MYSQL_DBNAME, RDS_PQ_SCHEMA
@@ -88,7 +90,12 @@ def get_sql_result_tool(profile, sql):
                 RDS_MYSQL_DBNAME=RDS_MYSQL_DBNAME,
             ))
         else:
-            engine = db.create_engine(p_db_url)
+            if profile['db_type'] == "bigquery":
+                password, host = ConnectionManagement.get_db_password_host_by_name(profile['conn_name'])
+                password = json.loads(password)
+                engine = db.create_engine(url=host, credentials_info=password)
+            else:
+                engine = db.create_engine(p_db_url)
         with engine.connect() as connection:
             logger.info(f'{sql=}')
             executed_result_df = pd.read_sql_query(text(sql), connection)