MSSQL changes to work also with Studio UI (#64)

peepdb/cli.py

@@ -1,10 +1,17 @@
 import click
-from .core import peep_db
-from .config import get_connection, save_connection, list_connections, remove_connection, remove_all_connections
 import json
+from .core import peep_db
+from .config import (
+    get_connection,
+    save_connection,
+    list_connections,
+    remove_connection,
+    remove_all_connections
+)
+from .exceptions import InvalidPassword
 from decimal import Decimal
 from datetime import date
-
+import logging
 
 class CustomEncoder(json.JSONEncoder):
     def default(self, obj):
@@ -22,7 +29,7 @@ def cli():
     peepDB: A quick database table viewer.
 
     This tool allows you to quickly inspect database tables without writing SQL queries.
-    It supports MySQL, PostgreSQL, MariaDB, SQLite, MongoDB and FireBase.
+    It supports MySQL, PostgreSQL, MariaDB, SQLite, MongoDB, Firebase, and MSSQL.
 
     Usage examples:
 
@@ -45,69 +52,34 @@ def cli():
 
 @cli.command()
 @click.argument('connection_name')
-@click.option('--table', help='Specific table to view')
-@click.option('--format', type=click.Choice(['table', 'json']), default='table', help='Output format')
-@click.option('--page', type=int, default=1, help='Page number for pagination')
-@click.option('--page-size', type=int, default=100, help='Number of rows per page')
-@click.option('--scientific', is_flag=True, help='Display numbers in scientific notation')
-def view(connection_name, table, format, page, page_size, scientific):
-    """
-    View database tables.
-
-    CONNECTION_NAME is the name of the saved database connection to use.
-
-    Examples:
-    peepdb view mydb
-    peepdb view mydb --table users --page 2 --page-size 50
-    peepdb view mydb --format json
-    """
-    connection = get_connection(connection_name)
-    if not connection:
-        click.echo(f"Error: No saved connection found with name '{connection_name}'.")
-        return
-
-    db_type, host, user, password, database = connection
-    result = peep_db(db_type, host, user, password, database, table, format=format, page=page, page_size=page_size, scientific=scientific)
-
-    if format == 'table':
-        click.echo(result)
-        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}")
-    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', 'sqlite', 'mongodb', 'firebase']), required=True, help='Database type')
+@click.option('--db-type',
+              type=click.Choice(['mysql', 'postgres', 'mariadb', 'sqlite', 'mongodb', 'firebase', 'mssql']),
+              required=True,
+              help='Database type')
 @click.option('--host', required=True, help='Database host or file path for SQLite')
+@click.option('--port', type=int, required=False, help='Database port (optional)')
 @click.option('--user', required=False, help='Database user (not required for SQLite)')
 @click.option('--password', required=False, help='Database password (not required for SQLite)')
-@click.option('--database', required=False, help='Database name (not required for SQLite/Firebase)')
-def save(connection_name, db_type, host, user, password, database):
+@click.option('--database', required=False, help='Database name (not required for SQLite/Firebase/MSSQL)')
+@click.option('--trusted', is_flag=True, help='Use Windows Authentication (Trusted Connection) for MSSQL')
+def save(connection_name, db_type, host, port, user, password, database, trusted):
     """
     Save a new database connection.
 
-    CONNECTION_NAME is the name to give to this saved connection.
-
     Example:
-    peepdb save postgres1 --db-type postgres --host localhost --user postgres --password YourPassword --database peepdb_test
+    peepdb save test_conn --db-type mssql --host localhost --port 1433 --trusted --database SomeDB
     """
-    if db_type in ['sqlite', 'firebase']:
-        user = password = database = ''  # Not used for SQLite and Firebase
-    else:
+    if trusted:
+        # For Windows Authentication, we skip user/password
+        user = ''
+        password = ''
+    elif db_type not in ['sqlite', 'firebase']:
         if not user:
             user = click.prompt('Enter username')
         if not password:
             password = click.prompt('Enter password', hide_input=True, confirmation_prompt=True)
-        if not database:
-            database = click.prompt('Enter database name')
-    save_connection(connection_name, db_type, host, user, password, database)
+
+    save_connection(connection_name, db_type, host, port, user, password, database, trusted)
     click.echo(f"Connection '{connection_name}' saved successfully.")
 
 
@@ -116,8 +88,6 @@ def list():
     """
     List saved connections.
 
-    This command displays all the database connections you have saved.
-
     Example:
     peepdb list
     """
@@ -131,8 +101,6 @@ def remove(connection_name):
     """
     Remove a specific saved connection.
 
-    CONNECTION_NAME is the name of the saved connection to remove.
-
     Example:
     peepdb remove mydb
     """
@@ -148,16 +116,64 @@ def remove_all():
     """
     Remove all saved connections.
 
-    This command will delete all your saved database connections.
-    Use with caution.
-
     Example:
     peepdb remove-all
     """
     count = remove_all_connections()
     click.echo(f"{count} connection(s) have been removed.")
 
 
+@cli.command()
+@click.argument('connection_name')
+@click.option('--table', required=False, help='Specific table to view')
+@click.option('--format', type=click.Choice(['table', 'json']), default='table', help='Output format')
+@click.option('--page', type=int, default=1, help='Page number to view')
+@click.option('--page-size', type=int, default=50, help='Number of rows per page')
+@click.option('--scientific', is_flag=True, help='Use scientific notation for numbers')
+def view(connection_name, table, format, page, page_size, scientific):
+    """
+    View data from a saved connection.
+
+    Example:
+    peepdb view mydb --table SomeTable
+    """
+    connection = get_connection(connection_name)
+    if not connection:
+        click.echo(f"Error: No saved connection found with name '{connection_name}'.")
+        return
+
+    # If there's a sixth item, it's 'trusted'
+    if len(connection) == 6:
+        db_type, host, user, password, database, trusted = connection
+    else:
+        db_type, host, user, password, database = connection
+        trusted = False
+
+    result = peep_db(
+        db_type=db_type,
+        host=host,
+        user=user,
+        password=password,
+        database=database,
+        table=table,
+        format=format,
+        page=page,
+        page_size=page_size,
+        scientific=scientific,
+        trusted=trusted
+    )
+
+    if format == 'table':
+        click.echo(result)
+        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}")
+    else:
+        click.echo(json.dumps(result, indent=2, cls=CustomEncoder))
+
+
 def main():
     cli()
 

peepdb/config.py

@@ -13,11 +13,10 @@
 from .exceptions import InvalidPassword
 
 @dataclass
-class KeySecurity():
+class KeySecurity:
     KEYRING = "os-keyring"
     PASSWORD = "password"
 
-
 CONFIG_DIR = os.path.expanduser("~/.peepdb")
 CONFIG_FILE = os.path.join(CONFIG_DIR, "config.json")
 SECURITY_CONFIG_FILE = os.path.join(CONFIG_DIR, "security_config.json")
@@ -26,10 +25,10 @@ class KeySecurity():
 
 logging.basicConfig(level=logging.DEBUG)
 logger = logging.getLogger(__name__)
+
 @cached(cache=TTLCache(maxsize=1024, ttl=600))
 def generate_key_from_password(salt):
-    password = click.prompt("Please entry the password to encrpyt/decrpty DB passwords",
-                            type=str).encode('utf-8')
+    password = click.prompt("Please enter the password to encrypt/decrypt DB passwords", type=str).encode('utf-8')
     salt = base64.b64decode(salt)
     kdf = PBKDF2HMAC(
         algorithm=hashes.SHA256(),
@@ -39,7 +38,6 @@ def generate_key_from_password(salt):
     )
     return base64.urlsafe_b64encode(kdf.derive(password)).decode("utf-8")
 
-
 @cached(cache=TTLCache(maxsize=1024, ttl=600))
 def fetch_key_from_keyring():
     key = keyring.get_password(KEYRING_SERVICE_NAME, KEYRING_USERNAME)
@@ -48,7 +46,6 @@ def fetch_key_from_keyring():
         keyring.set_password(KEYRING_SERVICE_NAME, KEYRING_USERNAME, key)
     return key
 
-
 def get_key_security_config():
     security_config = {}
     if os.path.exists(SECURITY_CONFIG_FILE):
@@ -60,29 +57,27 @@ def get_key_security_config():
 
     return security_config['PEEP_DB_KEY_SECURITY']
 
-
 def get_key():
     key_security_config = get_key_security_config()
     if key_security_config['type'] == KeySecurity.KEYRING:
         # Fetching encryption key from keyring
         key = fetch_key_from_keyring()
     else:
-        # Dynamically generating encryption key from users input
+        # Dynamically generating encryption key from user's input
         key = generate_key_from_password(key_security_config['salt'])
     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, port, user, password, database, trusted):
+    logger.debug(
+        f"Saving connection: {name}, {db_type}, {host}, {port}, {user}, {'*' * len(password) if password else 'None'}, {database}, Trusted: {trusted}"
+    )
 
-def save_connection(name, db_type, host, user, password, database):
-    logger.debug(f"Saving connection: {name}, {db_type}, {host}, {user}, {'*' * len(password) if password else 'None'}, {database}")
-
     if not os.path.exists(CONFIG_DIR):
         os.makedirs(CONFIG_DIR)
 
@@ -102,16 +97,17 @@ def save_connection(name, db_type, host, user, password, database):
         config[name] = {
             "db_type": db_type,
             "host": encrypt(host),
-            "user": encrypt(user),
-            "password": encrypt(password),
-            "database": encrypt(database)
+            "port": port,
+            "user": encrypt(user) if user else "",
+            "password": encrypt(password) if password else "",
+            "database": encrypt(database),
+            "trusted": trusted
         }
 
     with open(CONFIG_FILE, "w") as f:
         json.dump(config, f)
-
-    logger.debug("Connection saved successfully")
 
+    logger.debug("Connection saved successfully")
 
 def get_connection(name):
     if not os.path.exists(CONFIG_FILE):
@@ -128,23 +124,23 @@ def get_connection(name):
         if conn["db_type"] in ['sqlite', 'firebase']:
             return (
                 conn["db_type"],
-                conn["host"],  # For SQLite and Firebase, host is used directly
-                "",  # Empty string for user
-                "",  # Empty string for password
+                conn["host"],
+                "",  # user
+                "",  # password
                 conn.get("database", "")
             )
         else:
             return (
                 conn["db_type"],
                 decrypt(conn["host"]),
-                decrypt(conn["user"]),
-                decrypt(conn["password"]),
-                decrypt(conn["database"])
+                decrypt(conn["user"]) if conn["user"] else "",
+                decrypt(conn["password"]) if conn["password"] else "",
+                decrypt(conn["database"]),
+                conn.get("trusted", False)
             )
     except InvalidToken:
         raise InvalidPassword("Password is invalid !!!")
 
-
 def list_connections():
     if not os.path.exists(CONFIG_FILE):
         print("No saved connections.")
@@ -162,7 +158,6 @@ def list_connections():
         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
@@ -180,7 +175,6 @@ def remove_connection(name):
 
     return True
 
-
 def remove_all_connections():
     if not os.path.exists(CONFIG_FILE):
         return 0
@@ -193,19 +187,18 @@ def remove_all_connections():
 
         os.remove(CONFIG_FILE)
     except FileNotFoundError:
-        # File was deleted between check and remove
         pass
     except json.JSONDecodeError:
-        # File exists but is not valid JSON
         os.remove(CONFIG_FILE)
 
     return count
 
-
 def add_key_security():
     security_config = {}
-    key_security = click.prompt('Specify the way you want to store your encryption key',
-                                type=click.Choice([KeySecurity.KEYRING, KeySecurity.PASSWORD]))
+    key_security = click.prompt(
+        'Specify the way you want to store your encryption key',
+        type=click.Choice([KeySecurity.KEYRING, KeySecurity.PASSWORD])
+    )
     security_config['PEEP_DB_KEY_SECURITY'] = {"type": key_security}
     if key_security == KeySecurity.PASSWORD:
         security_config["PEEP_DB_KEY_SECURITY"]['salt'] = base64.b64encode(