From 114206097675bd439bce5e202d046ab9e80aa3c4 Mon Sep 17 00:00:00 2001 From: SangGyu An Date: Thu, 17 Aug 2023 23:03:38 -0700 Subject: [PATCH] improve error when calling sqlcmd (#804) --- CHANGELOG.md | 1 + src/sql/magic_cmd.py | 24 +++++++++++- src/tests/integration/test_questDB.py | 16 ++++---- src/tests/test_magic_cmd.py | 54 +++++++++++++++++++++++++++ 4 files changed, 85 insertions(+), 10 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4e6175e4d..6c2d747c9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,7 @@ * [Fix] Current connection and switching connections message only displayed when `feedback>=1` * [Fix] `--persist/--persist-replace` perform `ROLLBACK` automatically when needed * [Fix] `ResultSet` footer (when `displaylimit` truncates results and when showing how to convert to a data frame) now appears in the `ResultSet` plain text representation (#682) +* [Fix] Improve error when calling `%sqlcmd` (#761) * [API Change] When loading connections from a `.ini` file via `%sql --section section_name`, the section name is set as the connection alias * [API Change] Starting connections from a `.ini` file via `%sql [section_name]` has been deprecated * [Doc] Fixes documentation inaccuracy that said `:variable` was deprecated (we brought it back in `0.9.0`) diff --git a/src/sql/magic_cmd.py b/src/sql/magic_cmd.py index d816189b7..77bff8ef0 100644 --- a/src/sql/magic_cmd.py +++ b/src/sql/magic_cmd.py @@ -11,6 +11,7 @@ from sql.cmd.profile import profile from sql.cmd.explore import explore from sql.cmd.snippets import snippets +from sql.connection import ConnectionManager try: from traitlets.config.configurable import Configurable @@ -51,6 +52,14 @@ def _validate_execute_inputs(self, line): "explore", "snippets", ] + COMMANDS_CONNECTION_REQUIRED = [ + "tables", + "columns", + "test", + "profile", + "explore", + ] + COMMANDS_SQLALCHEMY_ONLY = ["tables", "columns", "test", "explore"] VALID_COMMANDS_MSG = ( f"Missing argument for %sqlcmd. " @@ -64,8 +73,19 @@ def _validate_execute_inputs(self, line): command, others = split[0].strip(), split[1:] if command in AVAILABLE_SQLCMD_COMMANDS: - if command not in ["profile"]: - util.support_only_sql_alchemy_connection("%sqlcmd") + if ( + command in COMMANDS_CONNECTION_REQUIRED + and not ConnectionManager.current + ): + raise exceptions.RuntimeError( + f"Cannot execute %sqlcmd {command} because there " + "is no active connection. Connect to a database " + "and try again." + ) + + if command in COMMANDS_SQLALCHEMY_ONLY: + util.support_only_sql_alchemy_connection(f"%sqlcmd {command}") + return self.execute(command, others) else: raise exceptions.UsageError( diff --git a/src/tests/integration/test_questDB.py b/src/tests/integration/test_questDB.py index 306306afe..cea75a899 100644 --- a/src/tests/integration/test_questDB.py +++ b/src/tests/integration/test_questDB.py @@ -615,17 +615,17 @@ def test_sql(ip_questdb, query, expected_results): @pytest.mark.parametrize( - "query", + "query, command", [ - ("%sqlcmd tables"), - ("%sqlcmd tables --schema some_schema"), - ("%sqlcmd columns --table penguins.csv"), - ("%sqlcmd test"), - ("%sqlcmd test --table penguins.csv"), + ("%sqlcmd tables", "tables"), + ("%sqlcmd tables --schema some_schema", "tables"), + ("%sqlcmd columns --table penguins.csv", "columns"), + ("%sqlcmd test", "test"), + ("%sqlcmd test --table penguins.csv", "test"), ], ) -def test_sqlcmd_not_supported_error(ip_questdb, query, capsys): - expected_error_message = f"%sqlcmd {NOT_SUPPORTED_SUFFIX}" +def test_sqlcmd_not_supported_error(ip_questdb, query, command, capsys): + expected_error_message = f"%sqlcmd {command} {NOT_SUPPORTED_SUFFIX}" with pytest.raises(UsageError) as excinfo: ip_questdb.run_cell(query) diff --git a/src/tests/test_magic_cmd.py b/src/tests/test_magic_cmd.py index c036b4e94..cd4cb3d6d 100644 --- a/src/tests/test_magic_cmd.py +++ b/src/tests/test_magic_cmd.py @@ -125,6 +125,60 @@ def test_error(tmp_empty, ip, cell, error_message): assert str(excinfo.value) == error_message +@pytest.mark.parametrize( + "command", + [ + "tables", + "columns", + "test", + "profile", + "explore", + ], +) +def test_sqlcmd_error_when_no_connection(ip_empty, command): + with pytest.raises(UsageError) as excinfo: + ip_empty.run_cell(f"%sqlcmd {command}") + + assert excinfo.value.error_type == "RuntimeError" + assert str(excinfo.value) == ( + f"Cannot execute %sqlcmd {command} because there is no " + "active connection. Connect to a database and try again." + ) + + +def test_sqlcmd_snippets_when_no_connection(ip_empty, capsys): + for key in list(store): + del store[key] + + ip_empty.run_cell("%sqlcmd snippets") + captured = capsys.readouterr() + assert "No snippets stored" in captured.out + + +@pytest.mark.parametrize( + "query, command", + [ + ("%sqlcmd tables", "tables"), + ("%sqlcmd columns --table penguins.csv", "columns"), + ( + "%sqlcmd test --table penguins.csv --column body_mass_g --greater 2900", + "test", + ), + ("%sqlcmd explore --table penguins.csv", "explore"), + ], +) +def test_sqlcmd_not_supported_error(ip_with_connections, query, command, capsys): + ip_with_connections.run_cell("%sql duckdb_dbapi") + expected_error_message = ( + f"%sqlcmd {command} is only supported with SQLAlchemy connections, " + "not with DBAPI connections" + ) + with pytest.raises(UsageError) as excinfo: + ip_with_connections.run_cell(query) + + assert expected_error_message in str(excinfo.value) + + def test_tables(ip): out = ip.run_cell("%sqlcmd tables").result._repr_html_() assert "author" in out