revert back to new head

app.py

@@ -9,7 +9,7 @@ import os
 import warnings
 
 
-from persistStorage import saveLog, getAllLogFilesPaths
+from persistStorage import saveLog, getAllLogFilesPaths, getNewCsvFilePath, removeAllCsvFiles
 from config import *
 from constants import *
 from utils import *
@@ -19,7 +19,7 @@ from queryHelperManagerCoT import QueryHelperChainOfThought
 
 
 pd.set_option('display.max_columns', None)
-pd.set_option('display.max_rows', 10)
+pd.set_option('display.max_rows', None)
 
 # Filter out all warning messages
 warnings.filterwarnings("ignore")
@@ -107,13 +107,40 @@ def preProcessSQL(sql):
   sql=sql.replace(';', '')
   disclaimerOutputStripping = ""
   if ('limit' in sql[-15:].lower())==False:
-    sql = sql + ' ' + 'limit 5'
-    disclaimerOutputStripping = """Results are stripped to show only top 5 rows.
-Please add your custom limit to get extend result.
-eg\n select * from schema.table limit 20\n\n"""
+    sql = sql + ' ' + 'limit 100'
+    disclaimerOutputStripping = """Results are stripped to show only top 100 rows.
+Please add your custom limit to get extended result.
+eg\n select * from schema.table limit 200n\n"""
   sql = sqlparse.format(sql, reindent=True, keyword_case='upper')
   return sql, disclaimerOutputStripping
 
+def onGetResultCsvFile(sql):
+  global dbEngine, queryHelper
+  sql, disclaimerOutputStripping = preProcessSQL(sql=sql)
+  if not isDataQuery(sql):
+    return "Sorry not allowed to run. As the query modifies the data."
+  try:
+    dbEngine2 = DbEngine(dbCreds)
+    dbEngine2.connect()
+    conn = dbEngine2.getConnection()
+    df = pd.read_sql_query(sql, con=conn)
+    
+    dbEngine2.disconnect()
+    # return disclaimerOutputStripping + str(pd.DataFrame(df))
+  except Exception as e:
+    # errorMessage = {"function":"testSQL","error":str(e), "userInput":sql}
+    # saveLog(errorMessage, 'error')
+    dbEngine2.disconnect()
+    df = pd.DataFrame()
+    # print(f"Error occured during running the query {sql}.\n and the error is {str(e)}")
+
+  removeAllCsvFiles()
+  csvFilePath = getNewCsvFilePath()
+  df.to_csv(csvFilePath, index=False)
+  downloadableFilesPaths = getAllLogFilesPaths()
+  fileComponent = gr.File(csvFilePath)
+  return fileComponent
+
 def testSQL(sql):
   global dbEngine, queryHelper
 
@@ -235,6 +262,11 @@ with gr.Blocks() as demo:
         text_button = gr.Button("RUN QUERY")
         clear = gr.ClearButton([text_input, text_output])
         text_button.click(testSQL, inputs=text_input, outputs=text_output)
+
+        csvFileComponent = gr.File([], file_count='multiple')
+        downloadCsv = gr.Button("Get result as csv")
+        downloadCsv.click(onGetResultCsvFile, inputs=text_input, outputs=csvFileComponent)
+        
      # screen 3 : To set creds, schema, tables and columns
     with gr.Tab("Setup"):
         gr.Markdown("""<h1><center> Run Query </center></h1>""")

configProd.py

@@ -22,7 +22,8 @@ PASSWD = os.getenv("passwd")
 
 DB_CREDS_DATA = ({"database":dbName, "user":userDB, "password":pwdDB, "host":host, "port":port})
 
+STORAGE_DIR = os.getenv("HF_HOME", "/data")
+logsDir = STORAGE_DIR
 
-logsDir = os.getenv("HF_HOME", "/data")
-
-TABLES_DATA_DIR = os.path.join(os.getenv("HF_HOME", "/data"), "tablesData")
+TABLES_DATA_DIR = os.path.join(STORAGE_DIR, "tablesData")
+RESULT_CSV_DIR = os.path.join(STORAGE_DIR, "csvResults")

persistStorage.py

@@ -3,20 +3,41 @@ import sqlite3
 from datetime import datetime, timedelta
 import pytz
 import os 
-from config import HUGGING_FACE_TOKEN, TABLES_DATA_DIR, logsDir
+from config import HUGGING_FACE_TOKEN, TABLES_DATA_DIR, logsDir, RESULT_CSV_DIR
 import pandas as pd
 import csv 
+import random
 
 try:
     os.makedirs(TABLES_DATA_DIR, exist_ok=True)
 except:
     pass
 
+try:
+    os.makedirs(RESULT_CSV_DIR, exist_ok=True)
+except:
+    pass
+
 # Set the time zone to Pacific Time Zone
 TIME_ZONE = 'US/Pacific'
 TIMEZONE_OBJ = pytz.timezone(TIME_ZONE)
 CACHE_TIME_EXPIRE = 20 #days
 
+def getNewCsvFilePath():
+    fileName = "ResultCsv_" + "".join([str(random.randint(0,9)) for i in range(3)]) + ".csv"
+    fileNameWithpath = os.path.join(RESULT_CSV_DIR, fileName)
+    return fileNameWithpath
+    
+def removeAllCsvFiles():
+    files = os.listdir(RESULT_CSV_DIR)
+    for fileName in files:
+        fileNameWithPath = os.path.join(RESULT_CSV_DIR, fileName)
+        try:
+            os.remove(fileNameWithPath)
+        except:
+            pass
+    
+
 def append_dict_to_csv(file_path, row_data):
     fieldnames = row_data.keys()