Code Restructure for implementing new features (#29)

* Addded Section for IDE Folders and Files

* Changes in Format to adhere to PEP 8 guidelines for Python code

* Extracted DB specific logic from core.py into independent files

.gitignore

@@ -14,6 +14,9 @@ build/
 # Environment variables
 .env
 
+# IDE
+.idea/
+
 # Operating system files
 .DS_Store
 Thumbs.db

peepdb/__main__.py

@@ -1,4 +1,4 @@
 from .cli import main
 
 if __name__ == "__main__":
-    main()
+    main()

peepdb/cli.py

@@ -5,6 +5,7 @@
 from decimal import Decimal
 from datetime import date
 
+
 class CustomEncoder(json.JSONEncoder):
     def default(self, obj):
         if isinstance(obj, Decimal):
@@ -13,6 +14,7 @@ def default(self, obj):
             return obj.isoformat()
         return super(CustomEncoder, self).default(obj)
 
+
 @click.group()
 @click.version_option()
 def cli():
@@ -33,6 +35,7 @@ def cli():
     """
     pass
 
+
 @cli.command()
 @click.argument('connection_name')
 @click.option('--table', help='Specific table to view')
@@ -63,11 +66,14 @@ def view(connection_name, table, format, page, page_size):
         if table:
             click.echo("\nNavigation:")
             click.echo(f"Current Page: {page}")
-            click.echo(f"Next Page: peepdb view {connection_name} --table {table} --page {page + 1} --page-size {page_size}")
-            click.echo(f"Previous Page: peepdb view {connection_name} --table {table} --page {max(1, page - 1)} --page-size {page_size}")
+            click.echo(
+                f"Next Page: peepdb view {connection_name} --table {table} --page {page + 1} --page-size {page_size}")
+            click.echo(
+                f"Previous Page: peepdb view {connection_name} --table {table} --page {max(1, page - 1)} --page-size {page_size}")
     else:
         click.echo(json.dumps(result, indent=2, cls=CustomEncoder))
 
+
 @cli.command()
 @click.argument('connection_name')
 @click.option('--db-type', type=click.Choice(['mysql', 'postgres', 'mariadb']), required=True, help='Database type')
@@ -87,6 +93,7 @@ def save(connection_name, db_type, host, user, password, database):
     save_connection(connection_name, db_type, host, user, password, database)
     click.echo(f"Connection '{connection_name}' saved successfully.")
 
+
 @cli.command()
 def list():
     """
@@ -99,6 +106,7 @@ def list():
     """
     list_connections()
 
+
 @cli.command()
 @click.argument('connection_name')
 @click.confirmation_option(prompt='Are you sure you want to remove this connection?')
@@ -116,6 +124,7 @@ def remove(connection_name):
     else:
         click.echo(f"No connection named '{connection_name}' found.")
 
+
 @cli.command()
 @click.confirmation_option(prompt='Are you sure you want to remove ALL saved connections?')
 def remove_all():
@@ -131,8 +140,10 @@ def remove_all():
     count = remove_all_connections()
     click.echo(f"{count} connection(s) have been removed.")
 
+
 def main():
     cli()
 
+
 if __name__ == '__main__':
-    main()
+    main()

peepdb/config.py

@@ -6,6 +6,7 @@
 CONFIG_FILE = os.path.join(CONFIG_DIR, "config.json")
 KEY_FILE = os.path.join(CONFIG_DIR, "key.key")
 
+
 def get_key():
     if not os.path.exists(KEY_FILE):
         key = Fernet.generate_key()
@@ -16,12 +17,15 @@ def get_key():
             key = key_file.read()
     return key
 
+
 def encrypt(message: str) -> str:
     return Fernet(get_key()).encrypt(message.encode()).decode()
 
+
 def decrypt(token: str) -> str:
     return Fernet(get_key()).decrypt(token.encode()).decode()
 
+
 def save_connection(name, db_type, host, user, password, database):
     if not os.path.exists(CONFIG_DIR):
         os.makedirs(CONFIG_DIR)
@@ -42,6 +46,7 @@ def save_connection(name, db_type, host, user, password, database):
     with open(CONFIG_FILE, "w") as f:
         json.dump(config, f)
 
+
 def get_connection(name):
     if not os.path.exists(CONFIG_FILE):
         return None
@@ -61,6 +66,7 @@ def get_connection(name):
         decrypt(conn["database"])
     )
 
+
 def list_connections():
     if not os.path.exists(CONFIG_FILE):
         print("No saved connections.")
@@ -77,7 +83,8 @@ def list_connections():
     for name, details in config.items():
         db_type = details.get('db_type', 'Unknown')
         print(f"- {name} ({db_type})")
-
+
+
 def remove_connection(name):
     if not os.path.exists(CONFIG_FILE):
         return False
@@ -95,6 +102,7 @@ def remove_connection(name):
 
     return True
 
+
 def remove_all_connections():
     if not os.path.exists(CONFIG_FILE):
         return 0
@@ -112,5 +120,5 @@ def remove_all_connections():
     except json.JSONDecodeError:
         # File exists but is not valid JSON
         os.remove(CONFIG_FILE)
-    
-    return count
+
+    return count

peepdb/core.py

@@ -1,3 +1,4 @@
+from peepdb.dbtypes import peepdb_mysql, peepdb_mariadb, peepdb_postgresql
 import mysql.connector
 import psycopg2
 import pymysql
@@ -10,44 +11,53 @@
 logging.basicConfig(level=logging.DEBUG)
 logger = logging.getLogger(__name__)
 
+
 def connect_to_database(db_type, host, user, password, database, port=None, **kwargs):
-       logger.debug(f"Attempting to connect to {db_type} database '{database}' on host '{host}' with user '{user}'")
-       try:
-           if db_type == 'mysql':
-               conn = mysql.connector.connect(host=host, user=user, password=password, database=database, port=port or 3306, **kwargs)
-           elif db_type == 'postgres':
-               conn = psycopg2.connect(host=host, user=user, password=password, database=database, port=port or 5432, **kwargs)
-           elif db_type == 'mariadb':
-               conn = pymysql.connect(host=host, user=user, password=password, database=database, port=port or 3306, **kwargs)
-           else:
-               raise ValueError("Unsupported database type")
-           logger.debug("Connection successful")
-           return conn
-       except (mysql.connector.Error, psycopg2.Error, pymysql.Error) as e:
-           logger.error(f"Failed to connect to {db_type} database: {str(e)}")
-           raise ConnectionError(f"Failed to connect to {db_type} database: {str(e)}")
-
+    logger.debug(f"Attempting to connect to {db_type} database '{database}' on host '{host}' with user '{user}'")
+    try:
+        if db_type == 'mysql':
+            conn = peepdb_mysql.connect_to_db(host=host, user=user, password=password, database=database,
+                                              port=port or 3306, **kwargs)
+        elif db_type == 'mariadb':
+            conn = peepdb_mariadb.connect_to_db(host=host, user=user, password=password, database=database,
+                                                port=port or 3306, **kwargs)
+        elif db_type == 'postgres':
+            conn = peepdb_postgresql.connect_to_db(host=host, user=user, password=password, database=database,
+                                                   port=port or 5432, **kwargs)
+        else:
+            raise ValueError("Unsupported database type")
+        logger.debug("Connection successful")
+        return conn
+    except (mysql.connector.Error, psycopg2.Error, pymysql.Error) as e:
+        logger.error(f"Failed to connect to {db_type} database: {str(e)}")
+        raise ConnectionError(f"Failed to connect to {db_type} database: {str(e)}")
+
+
 def fetch_tables(cursor, db_type):
-    if db_type in ['mysql', 'mariadb']:
-        cursor.execute("SHOW TABLES")
+    tables = []
+    if db_type == 'mysql':
+        tables = peepdb_mysql.fetch_tables(cursor)
+    elif db_type == 'mariadb':
+        tables = peepdb_mariadb.fetch_tables(cursor)
     elif db_type == 'postgres':
-        cursor.execute("SELECT table_name FROM information_schema.tables WHERE table_schema = 'public'")
-    return [table[0] for table in cursor.fetchall()]
+        tables = peepdb_postgresql.fetch_tables(cursor)
+    return [table[0] for table in tables]
+
 
 def view_table(cursor, table_name, page=1, page_size=100):
     # Get total number of rows
     cursor.execute(f"SELECT COUNT(*) FROM {table_name}")
     total_rows = cursor.fetchone()[0]
-    
+
     # Calculate total pages
     total_pages = math.ceil(total_rows / page_size)
-    
+
     # Ensure page is within bounds
     page = max(1, min(page, total_pages))
-    
+
     # Calculate offset
     offset = (page - 1) * page_size
-    
+
     # Fetch data for the current page
     cursor.execute(f"SELECT * FROM {table_name} LIMIT {page_size} OFFSET {offset}")
     columns = [col[0] for col in cursor.description]
@@ -61,14 +71,15 @@ def view_table(cursor, table_name, page=1, page_size=100):
                 value = value.isoformat()
             row_dict[columns[i]] = value
         rows.append(row_dict)
-    
+
     return {
         'data': rows,
         'page': page,
         'total_pages': total_pages,
         'total_rows': total_rows
     }
 
+
 def peep_db(db_type, host, user, password, database, table=None, format='table', page=1, page_size=100):
     conn = connect_to_database(db_type, host, user, password, database)
     cursor = conn.cursor()
@@ -87,6 +98,7 @@ def peep_db(db_type, host, user, password, database, table=None, format='table',
     else:
         return result
 
+
 def format_as_table(data):
     formatted_result = []
     for table_name, table_data in data.items():
@@ -97,6 +109,7 @@ def format_as_table(data):
             formatted_result.append(tabulate(table_rows, headers=headers, tablefmt='grid'))
         else:
             formatted_result.append("No data")
-        formatted_result.append(f"Page {table_data['page']} of {table_data['total_pages']} (Total rows: {table_data['total_rows']})")
+        formatted_result.append(
+            f"Page {table_data['page']} of {table_data['total_pages']} (Total rows: {table_data['total_rows']})")
         formatted_result.append("")  # Add an empty line between tables
-    return "\n".join(formatted_result).strip()
+    return "\n".join(formatted_result).strip()

peepdb/dbtypes/peepdb_mariadb.py

@@ -0,0 +1,38 @@
+import pymysql
+
+
+def connect_to_db(host, user, password, database, port, **kwargs):
+    """
+    Establishes a connection to database using the provided credentials.
+
+    Args:
+        host (str): The hostname or IP address of the server.
+        user (str): The username to use for the connection.
+        password (str): The password associated with the username.
+        database (str): The name of the database to connect to.
+        port (int): The port number on which the server is listening. Defaults to 3306 if not provided.
+
+    Returns:
+        pymysql.connections.Connection: A connection object to the database.
+
+    Raises:
+        pymysql.MySQLError
+    """
+    return pymysql.connect(host=host, user=user, password=password, database=database, port=port or 3306, **kwargs)
+
+
+def fetch_tables(cursor):
+    """
+    Retrieves the list of tables in the current database.
+
+    Args:
+        cursor (pymysql.cursors.Cursor): Cursor object used to execute SQL queries.
+
+    Returns:
+        list of tuple: A list of tuples where each tuple contains the name of a table in the database.
+
+    Raises:
+        pymysql.MySQLError
+    """
+    cursor.execute("SHOW TABLES")
+    return cursor.fetchall()

peepdb/dbtypes/peepdb_mysql.py

@@ -0,0 +1,40 @@
+import mysql.connector
+
+
+def connect_to_db(host, user, password, database, port, **kwargs):
+    """
+    Establishes a connection to database using the provided parameters.
+
+    Args:
+        host (str): The hostname or IP address of the server.
+        user (str): The username to use for the connection.
+        password (str): The password associated with the username.
+        database (str): The name of the database to connect to.
+        port (int): The port number on which the server is listening. Defaults to 3306 if not provided.
+
+    Returns:
+        mysql.connector.MySQLConnection: A MySQL connection object.
+
+    Raises:
+        mysql.connector.Error
+    """
+    return mysql.connector.connect(host=host, user=user, password=password, database=database, port=port or 3306,
+                                   **kwargs)
+
+
+def fetch_tables(cursor):
+    """
+    Retrieves the list of tables in the current database.
+
+    Args:
+        cursor (mysql.connector.cursor): Cursor object used to execute SQL queries.
+
+    Returns:
+        list of tuple: A list of tuples where each tuple contains the name of a table in the database.
+
+    Raises:
+        mysql.connector.Error
+    """
+    cursor.execute("SHOW TABLES")
+    return cursor.fetchall()
+