Code to add samples and execute query (#9)

Added:
1. Added prompt to ask user if they want to execute the query
2. A method to extract table names from the query generated
3. Added method to use PandaSQL to execute the query in memory sqllite

Author: narasimhard

requirements.txt

@@ -65,3 +65,4 @@ typing-inspect==0.9.0 ; python_full_version >= "3.8.16" and python_version < "3.
 urllib3==2.0.2 ; python_full_version >= "3.8.16" and python_version < "3.10"
 win32-setctime==1.1.0 ; python_full_version >= "3.8.16" and python_version < "3.10" and sys_platform == "win32"
 yarl==1.9.2 ; python_full_version >= "3.8.16" and python_version < "3.10"
+pandasql==0.7.3 ; python_full_version >= "3.8.16" and python_version < "3.10"

sidekick/db_config.py

@@ -1,10 +1,12 @@
 # create db with supplied info
 import json
 from pathlib import Path
+import pandas as pd
 
 import psycopg2 as pg
 import sqlalchemy
 from psycopg2.extras import Json
+from pandasql import sqldf
 from sidekick.logger import logger
 from sqlalchemy import create_engine
 from sqlalchemy_utils import database_exists
@@ -127,16 +129,61 @@ def has_table(self):
         )
         return sqlalchemy.inspect(engine).has_table(self.table_name)
 
-    def add_samples(self):
-        # Non-functional for now.
-        conn = pg.connect(
-            database=self.db_name, user=self.user_name, password=self.password, host=self.hostname, port=self.port
-        )
-        # Creating a cursor object using the cursor() method
-        conn.autocommit = True
-        cursor = conn.cursor()
+    def add_samples(self, data_csv_path=None):
+        conn_str = f"{self.dialect}://{self.user_name}:{self.password}@{self.hostname}:{self.port}/{self.db_name}"
+        try:
+            df = pd.read_csv(data_csv_path, infer_datetime_format=True)
+            engine = create_engine(conn_str, isolation_level='AUTOCOMMIT')
+
+            sample_query = f'SELECT COUNT(*) AS ROWS FROM {self.table_name} LIMIT 1'
+            num_rows_bef = pd.read_sql_query(sample_query, engine)
+
+            # Write rows to database
+            res = df.to_sql(self.table_name, engine, if_exists='append', index=False)
+
+            # Fetch the number of rows from the table
+            num_rows_aft = pd.read_sql_query(sample_query, engine)
+
+            logger.info(f"Number of rows inserted: {num_rows_aft.iloc[0, 0] - num_rows_bef.iloc[0, 0]}")
+
+            engine.dispose()
 
-        cursor.execute()
+        except Exception as e:
+            logger.info(f"Error occurred : {format(e)}")
+        finally:
+            engine.dispose()
 
-        # Commit your changes in the database
-        conn.commit()
+    def execute_query_db(self, query=None, n_rows=100):
+        try:
+            if query:
+                # Create an engine
+                conn_str = f"{self.dialect}://{self.user_name}:{self.password}@{self.hostname}:{self.port}/{self.db_name}"
+                engine = create_engine(conn_str)
+
+                # Create a connection
+                connection = engine.connect()
+
+                result = connection.execute(query)
+
+                # Process the query results
+                cnt = 0
+                logger.info("Here are the results from the queries: ")
+                for row in result:
+                    if cnt <= n_rows:
+                        # Access row data using row[column_name]
+                        logger.info(row)
+                        cnt += 1
+                    else:
+                        break
+                # Close the connection
+                connection.close()
+
+                # Close the engine
+                engine.dispose()
+            else:
+                logger.info("Query Empty or None!")
+        except Exception as e:
+            logger.info(f"Error occurred : {format(e)}")
+        finally:
+            connection.close()
+            engine.dispose()

sidekick/prompter.py

@@ -12,7 +12,7 @@
 from sidekick.db_config import DBConfig
 from sidekick.memory import EntityMemory
 from sidekick.query import SQLGenerator
-from sidekick.utils import save_query, setup_dir
+from sidekick.utils import save_query, setup_dir, extract_table_names, execute_query_pd
 
 # Load the config file and initialize required paths
 base_path = (Path(__file__).parent / "../").resolve()
@@ -51,6 +51,10 @@ def enter_table_name():
     val = input(color(F.GREEN, "", "Would you like to create a table for the database? (y/n): "))
     return val
 
+def enter_file_path(table: str):
+    val = input(color(F.GREEN, "", f"Please input the CSV file path to table: {table} : "))
+    return val
+
 
 @configure.command("log", help="Adjust log settings")
 @click.option("--set_level", "-l", help="Set log level (Default: INFO)")
@@ -162,9 +166,10 @@ def db_setup(db_name: str, hostname: str, user_name: str, password: str, port: i
         # Check if table exists; pending --> and doesn't have any rows
         if db_obj.has_table():
             click.echo(f"Checked table {db_obj.table_name} exists in the DB.")
-            val = input(color(F.GREEN, "", "Would you like to add few sample rows (at-least 3)? (y/n): "))
-            if val.lower() == "y":
-                db_obj.add_samples()
+            val = input(color(F.GREEN, "", "Would you like to add few sample rows (at-least 3)? (y/n):"))
+            if val.lower().strip() == "y" or val.lower().strip() == "yes":
+                val = input("Path to a CSV file to insert data from:")
+                db_obj.add_samples(val)
             else:
                 click.echo("Exiting...")
                 return
@@ -336,6 +341,44 @@ def query(question: str, table_info_path: str, sample_queries: str):
             _val = updated_sql if updated_sql else res
             save_query(base_path, query=question, response=_val)
 
+        exe_sql = click.prompt("Would you like to execute the generated SQL (y/n)?")
+        if exe_sql.lower() == "y" or exe_sql.lower() == "yes":
+            # For the time being, the default option is Pandas, but the user can be asked to select Database or Panadas DF later.
+            option = "pandas" # or DB
+            _val = updated_sql if updated_sql else res
+            if option == "DB":
+                hostname = env_settings["LOCAL_DB_CONFIG"]["HOST_NAME"]
+                user_name = env_settings["LOCAL_DB_CONFIG"]["USER_NAME"]
+                password = env_settings["LOCAL_DB_CONFIG"]["PASSWORD"]
+                port = env_settings["LOCAL_DB_CONFIG"]["PORT"]
+                db_name = env_settings["LOCAL_DB_CONFIG"]["DB_NAME"]
+
+                db_obj = DBConfig(db_name, hostname, user_name, password, port, base_path=base_path)
+                db_obj.execute_query(query=_val)
+            elif option == "pandas":
+                tables = extract_table_names(_val)
+                tables_path = dict()
+                for table in tables:
+                    while True:
+                        val = enter_file_path(table)
+                        if not os.path.isfile(val):
+                            click.echo("In-correct Path. Please enter again! Yes(y) or no(n)")
+                            # val = enter_file_path(table)
+                        else:
+                            tables_path[table] = val
+                            break
+
+                assert len(tables) == len(tables_path)
+
+                res = execute_query_pd(query=_val, tables_path=tables_path, n_rows=100)
+
+                logger.info("The query results are:")
+                logger.info(res)
+
+            else:
+                click.echo("Exiting...")
+
+
 
 if __name__ == "__main__":
     cli()

sidekick/utils.py

@@ -5,6 +5,8 @@
 
 import numpy as np
 import pandas as pd
+from pandasql import sqldf
+import re
 from sentence_transformers import SentenceTransformer
 from sidekick.logger import logger
 from sklearn.metrics.pairwise import cosine_similarity
@@ -106,3 +108,42 @@ def csv_parser(input_path: str):
     # ]
     res = df.apply(lambda row: f"# query: {row['query']}\n# answer: {row['answer']}", axis=1).to_list()
     return res
+
+def extract_table_names(query: str):
+    """
+    Extracts table names from a SQL query.
+
+    Parameters:
+        query (str): The SQL query to extract table names from.
+
+    Returns:
+        list: A list of table names.
+    """
+    table_names = re.findall(r'\bFROM\s+(\w+)', query, re.IGNORECASE)
+    table_names += re.findall(r'\bJOIN\s+(\w+)', query, re.IGNORECASE)
+    table_names += re.findall(r'\bUPDATE\s+(\w+)', query, re.IGNORECASE)
+    table_names += re.findall(r'\bINTO\s+(\w+)', query, re.IGNORECASE)
+
+    # Below keywords may not be relevant for the project but adding for sake for completness
+    table_names += re.findall(r'\bINSERT\s+INTO\s+(\w+)', query, re.IGNORECASE)
+    table_names += re.findall(r'\bDELETE\s+FROM\s+(\w+)', query, re.IGNORECASE)
+
+    return table_names
+
+def execute_query_pd(query=None, tables_path=None, n_rows=100):
+    """
+    Runs an SQL query on a pandas DataFrame.
+
+    Parameters:
+        df (pandas DataFrame): The DataFrame to query.
+        query (str): The SQL query to execute.
+
+    Returns:
+        pandas DataFrame: The result of the SQL query.
+    """
+    for table in tables_path:
+          locals()[f"{table}"] = pd.read_csv(tables_path[table])
+
+    res_df = sqldf(query, locals())
+
+    return res_df