openai-sql/nogui/main.py
2023-05-04 10:27:54 +02:00

124 lines
4.2 KiB
Python

import getpass
import modules.aisql as aisql
class Main:
def __init__(self):
self.tableschema = {}
self.dbtype = 'postgresql'
self.ip = 'localhost'
self.port = '32768'
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('Configure the Database')
print(f'Database type')
print("postgresql - PostgreSQL")
print("mysql - MySQL (WIP)")
self.dbtype = input(f'Database type (current: {self.dbtype}): ') or self.dbtype
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
print('___')
if self.dbtype == 'postgresql':
# TODO: Add more Database languages
import database.postgresql as pg
self.db = pg.Postgres(self.ip, self.port, self.username, self.password, self.database)
return self.db
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. Get Tables Schema')
print('5. Exit')
option = input('Option: ')
print('___')
if option == '1':
self.generate_sql()
elif option == '2':
query = input('Query: ')
self.run_sql(query)
elif option == '3':
self.db = self.config()
elif option == '4':
if self.db is not None:
self.tableschema = self.db.get_schema()
print(self.tableschema)
else:
self.db = self.config()
self.tableschema = self.db.get_schema()
print(self.tableschema)
elif option == '5':
exit()
else:
print('Invalid option')
def generate_sql(self):
ai = aisql.AI(self.apikey)
generate = True
while generate:
text = input('Text: ')
print('___')
if self.db is None:
self.db = self.config()
sql = ai.humantosql(text, self.dbtype, self.db.get_schema())
print(f'SQL:\n{sql}\n___')
print('Would you like to run it?')
print('1. Yes')
print('2. No')
print('3. Edit')
choice = input('Choice: ')
print('___')
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):
if self.db is None:
self.db = self.config()
ai = aisql.AI(self.apikey)
function = ai.decide(query)
print(f'Using {function}')
try:
if "fetchall".casefold() in function.casefold():
print(self.db.fetchall(query))
elif "fetchone".casefold() in function.casefold():
print(self.db.fetchone(query))
elif "fetchmany".casefold() in function.casefold():
size = function.split('=')[1].strip(']')
print(self.db.fetchmany(query, size))
elif "execute".casefold() in function.casefold():
self.db.execute(query)
elif "executemany".casefold() in function.casefold():
size = function.split('=')[1].strip(']')
self.db.executemany(query, size)
else:
print('Invalid option')
print('___')
except Exception as e:
print('Something went wrong:')
print(e)
print('___')
if __name__ == '__main__':
main = Main()
while True:
main.select()