From ca09ec3b5b95964467c6b46354205924a7197527 Mon Sep 17 00:00:00 2001 From: crennis Date: Wed, 3 May 2023 16:28:37 +0200 Subject: [PATCH] Made a new Prototype --- nogui/database/postgresql.py | 80 +++++++++-------------- nogui/database/postgresqlignore.txt | 69 ++++++++++++++++++++ nogui/main.py | 99 +++++++++++++++++++++++++++++ nogui/modules/aisql.py | 32 ++++++++++ nogui/test.py | 24 +++++++ 5 files changed, 253 insertions(+), 51 deletions(-) create mode 100644 nogui/database/postgresqlignore.txt create mode 100644 nogui/main.py create mode 100644 nogui/modules/aisql.py create mode 100644 nogui/test.py diff --git a/nogui/database/postgresql.py b/nogui/database/postgresql.py index a03a601..8b42c62 100644 --- a/nogui/database/postgresql.py +++ b/nogui/database/postgresql.py @@ -1,90 +1,68 @@ import psycopg2 +import os class Postgres: - def __init__(self): - self.db_ip = 'localhost' - self.db_port = '32771' - self.db_username = 'root' - self.db_password = '' - self.db_name = '' - - def config(self, ip, port, username, password, name): + def __init__(self, ip, port, username, password, database): self.db_ip = ip self.db_port = port self.db_username = username self.db_password = password - self.db_name = name + self.db_name = database - def fetchall(self, query): - conn = psycopg2.connect( + self.conn = psycopg2.connect( host=self.db_ip, port=self.db_port, user=self.db_username, password=self.db_password, database=self.db_name ) - cur = conn.cursor() + self.tables = [] + fetch = self.fetchall("SELECT table_name FROM information_schema.tables WHERE table_schema = 'public';") + for item in fetch: + self.tables.append(item[0]) + + self.tableschema = {} + for table in self.tables: + fetch = self.fetchall(f"SELECT column_name, data_type FROM information_schema.columns WHERE table_name = '{table}';") + self.tableschema.update({table: fetch}) + + def __str__(self): + return f'PostgreSQL Server: {self.db_ip}:{self.db_port} as {self.db_username} on database {self.db_name}' + + def fetchall(self, query): + cur = self.conn.cursor() cur.execute(query) result = cur.fetchall() cur.close() - conn.close() return result def fetchone(self, query): - conn = psycopg2.connect( - host=self.db_ip, - port=self.db_port, - user=self.db_username, - password=self.db_password, - database=self.db_name - ) - cur = conn.cursor() + + cur = self.conn.cursor() cur.execute(query) result = cur.fetchone() cur.close() - conn.close() return result def fetchmany(self, query, size): - conn = psycopg2.connect( - host=self.db_ip, - port=self.db_port, - user=self.db_username, - password=self.db_password, - database=self.db_name - ) - cur = conn.cursor() + cur = self.conn.cursor() cur.execute(query) result = cur.fetchmany(size) cur.close() - conn.close() return result def execute(self, query): - conn = psycopg2.connect( - host=self.db_ip, - port=self.db_port, - user=self.db_username, - password=self.db_password, - database=self.db_name - ) - cur = conn.cursor() + cur = self.conn.cursor() cur.execute(query) - conn.commit() + self.conn.commit() cur.close() - conn.close() def executemany(self, query, values): - conn = psycopg2.connect( - host=self.db_ip, - port=self.db_port, - user=self.db_username, - password=self.db_password, - database=self.db_name - ) - cur = conn.cursor() + cur = self.conn.cursor() cur.executemany(query, values) - conn.commit() + self.conn.commit() cur.close() - conn.close() + + def closeconnection(self): + self.conn.close() diff --git a/nogui/database/postgresqlignore.txt b/nogui/database/postgresqlignore.txt new file mode 100644 index 0000000..bb4ff79 --- /dev/null +++ b/nogui/database/postgresqlignore.txt @@ -0,0 +1,69 @@ +"role_column_grants" +"information_schema_catalog_name" +"column_domain_usage" +"applicable_roles" +"administrable_role_authorizations" +"domain_constraints" +"attributes" +"column_privileges" +"character_sets" +"check_constraint_routine_usage" +"check_constraints" +"column_udt_usage" +"collations" +"collation_character_set_applicability" +"key_column_usage" +"column_column_usage" +"columns" +"domain_udt_usage" +"constraint_column_usage" +"constraint_table_usage" +"domains" +"referential_constraints" +"enabled_roles" +"parameters" +"routine_column_usage" +"routine_privileges" +"role_routine_grants" +"routine_routine_usage" +"table_privileges" +"routine_sequence_usage" +"routine_table_usage" +"udt_privileges" +"routines" +"role_table_grants" +"schemata" +"sequences" +"sql_features" +"tables" +"sql_implementation_info" +"sql_parts" +"sql_sizing" +"transforms" +"table_constraints" +"view_routine_usage" +"role_udt_grants" +"triggered_update_columns" +"triggers" +"user_defined_types" +"usage_privileges" +"role_usage_grants" +"view_column_usage" +"view_table_usage" +"views" +"data_type_privileges" +"_pg_foreign_table_columns" +"element_types" +"column_options" +"_pg_foreign_data_wrappers" +"foreign_data_wrapper_options" +"foreign_data_wrappers" +"_pg_foreign_servers" +"foreign_server_options" +"foreign_servers" +"_pg_foreign_tables" +"foreign_table_options" +"foreign_tables" +"_pg_user_mappings" +"user_mapping_options" +"user_mappings" \ No newline at end of file diff --git a/nogui/main.py b/nogui/main.py new file mode 100644 index 0000000..9a64414 --- /dev/null +++ b/nogui/main.py @@ -0,0 +1,99 @@ +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() \ No newline at end of file diff --git a/nogui/modules/aisql.py b/nogui/modules/aisql.py new file mode 100644 index 0000000..ef3fa03 --- /dev/null +++ b/nogui/modules/aisql.py @@ -0,0 +1,32 @@ +import openai + + +class AI: + def __init__(self, api_key, tableschema): + self.api_key = api_key + openai.api_key = self.api_key + self.convertlog = [{"role": "system", "content": f"You convert human language to SQL. You do not add any adidtional information. You are indirectly connected to a Database, so you can Query, Execute etc, just make the SQL statement. For better context you get the tables and columns from the database: {tableschema}"}] + + def humantosql(self, text) -> str: + prompt = {"role": "user", "content": text} + self.convertlog.append(prompt) + response = openai.ChatCompletion.create( + model="gpt-4", + messages=self.convertlog + ) + response = response['choices'][0]['message']['content'] + self.convertlog.append({"role": "system", "content": response}) + return response + + def decide(self, sql) -> str: + prompt = [{"role": "system", "content": "You have to decide which function it should use. Answer with [FETCHALL] to fetch all, [FETCHONE] to fetch only one, [FETCHMANY=N] to fetchmany with N being the range, [EXECUTE] to just execute, [EXECUTEMANY=N] to execute many with N being the range"}, + {"role": "user", "content": sql}] + + response = openai.ChatCompletion.create( + model="gpt-4", + messages=prompt + ) + + response = response['choices'][0]['message']['content'] + + return response diff --git a/nogui/test.py b/nogui/test.py new file mode 100644 index 0000000..a5b06d0 --- /dev/null +++ b/nogui/test.py @@ -0,0 +1,24 @@ +import database.postgresql as pg + +ip = 'localhost' +port = '32771' +username = 'root' +password = 'wv6G*Uny#4^CUA9E' +database = 'discord' + +default_tables = [''] + +postgres = pg.Postgres() +postgres.config(ip, port, username, password, database) + +tables = postgres.fetchall('SELECT table_name FROM information_schema.tables;') + +for table in tables: + if 'pg_' in table[0]: + continue + else: + print(table[0]) + columns = postgres.fetchall(f'SELECT column_name FROM information_schema.columns WHERE table_name = \'{table[0]}\';') + for column in columns: + print(column[0]) + print('\n') \ No newline at end of file