518 refactor magic cmd

CHANGELOG.md

@@ -1,12 +1,12 @@
 # CHANGELOG
 
 ## 0.7.10dev
-* [Feature] Support flexible spacing `myvar=<<` operator ([#525](https://github.com/ploomber/jupysql/issues/525))
 
+* [Fix] Refactored `magic_cmd.py` by assigning each command logic to a separate file [#518]
+* [Feature] Support flexible spacing `myvar=<<` operator ([#525](https://github.com/ploomber/jupysql/issues/525))
 * [Doc] Modified integrations content to ensure they're all consistent (#523)
 * [Doc] Document --persist-replace in API section (#539)
 * [Fix] Fixed CI issue by updating `invalid_connection_string_duckdb` in `test_magic.py` (#631)
-
 * [Fix] Refactored `ResultSet` to lazy loading (#470)
 
 ## 0.7.9 (2023-06-19)

src/sql/cmd/cmd_utils.py

@@ -0,0 +1,12 @@
+import argparse
+import sys
+from sql import exceptions
+
+
+class CmdParser(argparse.ArgumentParser):
+    def exit(self, status=0, message=None):
+        if message:
+            self._print_message(message, sys.stderr)
+
+    def error(self, message):
+        raise exceptions.UsageError(message)

src/sql/cmd/columns.py

@@ -0,0 +1,17 @@
+from sql import inspect
+from sql.util import sanitize_identifier
+from sql.cmd.cmd_utils import CmdParser
+
+
+def execute_columns_command(others):
+    """
+    Execution logic for the columns command
+
+    """
+    parser = CmdParser()
+
+    parser.add_argument("-t", "--table", type=str, help="Table name", required=True)
+    parser.add_argument("-s", "--schema", type=str, help="Schema name", required=False)
+
+    args = parser.parse_args(others)
+    return inspect.get_columns(name=sanitize_identifier(args.table), schema=args.schema)

src/sql/cmd/explore.py

@@ -0,0 +1,17 @@
+from sql.widgets import TableWidget
+from IPython.display import display
+
+from sql.cmd.cmd_utils import CmdParser
+
+
+def execute_expolore_command(others):
+    """
+    Execution logic for the explore command
+
+    """
+    parser = CmdParser()
+    parser.add_argument("-t", "--table", type=str, help="Table name", required=True)
+    args = parser.parse_args(others)
+
+    table_widget = TableWidget(args.table)
+    display(table_widget)

src/sql/cmd/profile.py

@@ -0,0 +1,27 @@
+from sql import inspect
+from sql.cmd.cmd_utils import CmdParser
+
+
+def execute_profile_command(others):
+    """
+    Execution logic for the profile command
+
+    """
+    parser = CmdParser()
+    parser.add_argument("-t", "--table", type=str, help="Table name", required=True)
+
+    parser.add_argument("-s", "--schema", type=str, help="Schema name", required=False)
+
+    parser.add_argument(
+        "-o", "--output", type=str, help="Store report location", required=False
+    )
+
+    args = parser.parse_args(others)
+
+    report = inspect.get_table_statistics(schema=args.schema, name=args.table)
+
+    if args.output:
+        with open(args.output, "w") as f:
+            f.write(report._repr_html_())
+
+    return report

src/sql/cmd/snippets.py

@@ -0,0 +1,9 @@
+from sql.sqlcmd import sqlcmd_snippets  # noqa
+
+
+def execute_snippets_command(others):
+    """
+    Execution logic for the snippets command
+
+    """
+    return sqlcmd_snippets(others)

src/sql/cmd/tables.py

@@ -0,0 +1,16 @@
+from sql import inspect
+from sql.cmd.cmd_utils import CmdParser
+
+
+def execute_tables_command(others):
+    """
+    Execution logic for the tables command
+
+    """
+    parser = CmdParser()
+
+    parser.add_argument("-s", "--schema", type=str, help="Schema name", required=False)
+
+    args = parser.parse_args(others)
+
+    return inspect.get_table_names(schema=args.schema)

src/sql/cmd/test.py

@@ -0,0 +1,162 @@
+from sql import exceptions
+import sql.connection
+from sqlalchemy import text
+from sqlglot import select, condition
+from prettytable import PrettyTable
+from sql.cmd.cmd_utils import CmdParser
+
+
+def return_test_results(args, conn, query):
+    try:
+        columns = []
+        column_data = conn.execute(text(query)).cursor.description
+        res = conn.execute(text(query)).fetchall()
+        for column in column_data:
+            columns.append(column[0])
+        res = [columns, *res]
+        return res
+    except Exception as e:
+        if "column" in str(e):
+            raise exceptions.UsageError(
+                f"Referenced column '{args.column}' not found!"
+            ) from e
+
+
+def run_each_individually(args, conn):
+    base_query = select("*").from_(args.table)
+
+    storage = {}
+
+    if args.greater:
+        where = condition(args.column + "<=" + args.greater)
+        current_query = base_query.where(where).sql()
+
+        res = return_test_results(args, conn, query=current_query)
+
+        if res is not None:
+            storage["greater"] = res
+    if args.greater_or_equal:
+        where = condition(args.column + "<" + args.greater_or_equal)
+
+        current_query = base_query.where(where).sql()
+
+        res = return_test_results(args, conn, query=current_query)
+
+        if res is not None:
+            storage["greater_or_equal"] = res
+
+    if args.less_than_or_equal:
+        where = condition(args.column + ">" + args.less_than_or_equal)
+        current_query = base_query.where(where).sql()
+
+        res = return_test_results(args, conn, query=current_query)
+
+        if res is not None:
+            storage["less_than_or_equal"] = res
+    if args.less_than:
+        where = condition(args.column + ">=" + args.less_than)
+        current_query = base_query.where(where).sql()
+
+        res = return_test_results(args, conn, query=current_query)
+
+        if res is not None:
+            storage["less_than"] = res
+    if args.no_nulls:
+        where = condition("{} is NULL".format(args.column))
+        current_query = base_query.where(where).sql()
+
+        res = return_test_results(args, conn, query=current_query)
+
+        if res is not None:
+            storage["null"] = res
+
+    return storage
+
+
+def execute_test_command(others):
+    """
+    Execution logic for the test command
+
+    """
+    parser = CmdParser()
+
+    parser.add_argument("-t", "--table", type=str, help="Table name", required=True)
+    parser.add_argument("-c", "--column", type=str, help="Column name", required=False)
+    parser.add_argument(
+        "-g",
+        "--greater",
+        type=str,
+        help="Greater than a certain number.",
+        required=False,
+    )
+    parser.add_argument(
+        "-goe",
+        "--greater-or-equal",
+        type=str,
+        help="Greater or equal than a certain number.",
+        required=False,
+    )
+    parser.add_argument(
+        "-l",
+        "--less-than",
+        type=str,
+        help="Less than a certain number.",
+        required=False,
+    )
+    parser.add_argument(
+        "-loe",
+        "--less-than-or-equal",
+        type=str,
+        help="Less than or equal to a certain number.",
+        required=False,
+    )
+    parser.add_argument(
+        "-nn",
+        "--no-nulls",
+        help="Returns rows in specified column that are not null.",
+        action="store_true",
+    )
+
+    args = parser.parse_args(others)
+
+    COMPARATOR_ARGS = [
+        args.greater,
+        args.greater_or_equal,
+        args.less_than,
+        args.less_than_or_equal,
+    ]
+
+    if args.table and not any(COMPARATOR_ARGS):
+        raise exceptions.UsageError("Please use a valid comparator.")
+
+    if args.table and any(COMPARATOR_ARGS) and not args.column:
+        raise exceptions.UsageError("Please pass a column to test.")
+
+    if args.greater and args.greater_or_equal:
+        return exceptions.UsageError(
+            "You cannot use both greater and greater "
+            "than or equal to arguments at the same time."
+        )
+    elif args.less_than and args.less_than_or_equal:
+        return exceptions.UsageError(
+            "You cannot use both less and less than "
+            "or equal to arguments at the same time."
+        )
+
+    conn = sql.connection.Connection.current.session
+    result_dict = run_each_individually(args, conn)
+
+    if any(len(rows) > 1 for rows in list(result_dict.values())):
+        for comparator, rows in result_dict.items():
+            if len(rows) > 1:
+                print(f"\n{comparator}:\n")
+                _pretty = PrettyTable()
+                _pretty.field_names = rows[0]
+                for row in rows[1:]:
+                    _pretty.add_row(row)
+                print(_pretty)
+        raise exceptions.UsageError(
+            "The above values do not not match your test requirements."
+        )
+    else:
+        return True