openai-sql/nogui/main.py
2023-05-03 16:28:37 +02:00

99 lines
3.3 KiB
Python

import getpass
import modules.aisql as aisql
class Main:
def __init__(self):
self.tableschema = None
self.dbtype = 'postgresql'
self.ip = 'localhost'
self.port = '32771'
self.username = 'root'
self.password = 'wv6G*Uny#4^CUA9E'
self.database = 'discord'
self.db = None
self.apikey = open('apikey', 'r').read().strip('\n')
def config(self) -> object:
print(f'Database type')
print("postgresql - PostgreSQL")
print("mysql - MySQL (WIP)")
self.dbtype = input('Database type: ')
self.ip = input(f'IP: (current: {self.ip})') or self.ip
self.port = input(f'Port: (current: {self.port})') or self.port
self.username = input(f'Username: (current: {self.username})') or self.username
self.password = getpass.getpass(f'Password: (current: {self.password})') or self.password
self.database = input(f'Database: (current: {self.database})') or self.database
if self.dbtype == 'postgresql':
# TODO: Add more Database languages
import database.postgresql as pg
db = pg.Postgres(self.ip, self.port, self.username, self.password, self.database)
return db, db.tableschema
def select(self):
print('Select an option to continue')
print('1. Convert human language to SQL')
print('2. Run SQL query')
print('3. Configure database')
print('4. Exit')
option = input('Option: ')
if option == '1':
self.generate_sql()
elif option == '2':
query = input('Query: ')
print(self.db.fetchall(query))
elif option == '3':
self.db, self.tableschema = self.config()
elif option == '4':
exit()
else:
print('Invalid option')
def generate_sql(self):
ai = aisql.AI(self.apikey, self.tableschema)
generate = True
while generate:
text = input('Text: ')
sql = ai.humantosql(text)
print(f'SQL:\n{sql}')
print('Would you like to run it?')
print('1. Yes')
print('2. No')
print('3. Edit')
choice = input('Choice: ')
if choice == '1':
self.run_sql(sql)
generate = False
elif choice == '2':
generate = False
elif choice == '3':
pass
else:
print('Invalid option')
def run_sql(self, query):
ai = aisql.AI(self.apikey, self.tableschema)
function = ai.decide(query)
print(f'Using {function}')
if "fetchall".lower() in function:
self.db.fetchall(query)
elif "fetchone".lower() in function:
self.db.fetchone(query)
elif "fetchmany".lower() in function:
size = function.split('=')[1].strip(']')
self.db.fetchmany(query, size)
elif "execute".lower() in function:
self.db.execute(query)
elif "executemany".lower() in function:
size = function.split('=')[1].strip(']')
self.db.executemany(query, size)
else:
print('Invalid option')
if __name__ == '__main__':
main = Main()
while True:
main.select()