From f3a877a7ffce2aeab8885bf22a8d86ac10ca9dae Mon Sep 17 00:00:00 2001 From: Roland Walker Date: Sat, 1 Aug 2026 17:04:36 -0400 Subject: [PATCH] let favorite queries use Jinja2 templates allowing optional arguments and control blocks in favorite queries. Include completions for Jinja2 key names. Supporting the completions is quite a bit more code than the core functionality! Completions also can only be offered if the key name is referred to explicitly in the template (not merely iterated over, which is technically possible). Extensively update the config file and internal helpdoc with examples and notes. It seems arbitrary to choose whether Jinja2 templates are processed before or after the legacy positional arguments are substituted. I chose Jinja2 before positional. The agent hacked together a UUID-based solution for this, which may not be pretty. There is one regression: whereas a positional argument could previously freely begin with the dash character, that can now be ambiguous with the --key=value syntax. The solution is to disambiguate by putting such positional values after a "--": /f query --key=value -- --ambiguous-positional-value-- This should be rare enough that the regression is worth the new feature. This adds a Jinja2 direct dependency to the project. It was already a transitive dependency via altair if mycli was installed with dataframe support. --- changelog.md | 5 + mycli/myclirc | 24 +++ mycli/packages/completion_engine.py | 52 +++++- mycli/packages/special/favoritequeries.py | 82 ++++++++- mycli/packages/special/iocommands.py | 114 +++++++++++- mycli/sqlcompleter.py | 28 ++- pyproject.toml | 1 + test/features/fixture_data/help_commands.txt | 2 +- test/features/named_queries.feature | 10 + test/features/steps/named_queries.py | 25 +++ test/myclirc | 24 +++ test/pytests/test_completion_engine.py | 23 +++ test/pytests/test_special_iocommands.py | 181 +++++++++++++++++++ test/pytests/test_sqlcompleter.py | 44 +++++ 14 files changed, 597 insertions(+), 18 deletions(-) diff --git a/changelog.md b/changelog.md index dec06a548..7109fb35f 100644 --- a/changelog.md +++ b/changelog.md @@ -1,6 +1,11 @@ Upcoming (TBD) ============== +Features +--------- +* Let favorite queries use Jinja2 templates, allowing optional arguments. + + Documentation --------- * Document shell completions in `README.md`. diff --git a/mycli/myclirc b/mycli/myclirc index 038dbf4d3..c60c9cad6 100644 --- a/mycli/myclirc +++ b/mycli/myclirc @@ -495,8 +495,32 @@ matching-bracket.other = '#000000 bg:#aacccc' # Favorite queries. # You can add your favorite queries here. They will be available in the # REPL when you type `/f` or `/f `. +# Define multiine favorite queries using '''triplequotes'''. [favorite_queries] # example = "SELECT * FROM example_table WHERE id = 1" +# +# Favorite queries can use the Jinja2 templating language. Named argument +# values are available as keys in a dictionary named "kv". +# Run this example templated query with: /f find_user --name=henry +# Any keys refereed to by name will be available as completions. +# find_user = "SELECT * FROM users WHERE name = '{{ kv.name }}'" +# +# Jinja2 parameters are optional. +# Run this example templated query with: /f recent_user --name=henry +# or alternatively, with different behavior: /f recent_user +# recent_user = "SELECT * FROM users WHERE create_date >= NOW() - INTERVAL 1 DAY {% if kv.name %} AND name = '{{ kv.name }}' {% endif %}" +# +# Since "kv" is a Python dictionary, collisions with existing dictionary +# methods will require extra syntax, as will keys with internal dashes: +# example = "SELECT * FROM example_table WHERE id = {{ kv['get'] }}" +# example = "SELECT * FROM example_table WHERE id = {{ kv['my-id'] }}" +# +# Mandatory positional arguments are also supported: +# Run this example templated query with: /f find_user henry +# find_user = "SELECT * FROM users WHERE name = '$1'" +# +# Plase ambiguous positional arguments after "--": +# /f find_user -- --unusual-name-- # Initial commands to execute when connecting to any database. [init-commands] diff --git a/mycli/packages/completion_engine.py b/mycli/packages/completion_engine.py index 3e90d8715..8562eec5a 100644 --- a/mycli/packages/completion_engine.py +++ b/mycli/packages/completion_engine.py @@ -1,6 +1,7 @@ from dataclasses import dataclass import functools import re +import shlex from typing import Any, Callable, Literal import sqlparse @@ -791,7 +792,10 @@ def suggest_special(text: str) -> list[dict[str, Any]]: if cmd.lower() in ('tableformat', '/tableformat', 'redirectformat', '/redirectformat'): return [{"type": "table_format"}] - if cmd in ["\\f", "/f", "\\fs", "/fs", "\\fd", "/fd"]: + if cmd in ["\\f", "/f"]: + return suggest_favorite_query_with_template(text, _arg) + + if cmd in ["\\fs", "/fs", "\\fd", "/fd"]: return [{"type": "favoritequery"}] if cmd in ["\\dt", "/dt", "\\dt+", "/dt+"]: @@ -870,6 +874,52 @@ def suggest_special(text: str) -> list[dict[str, Any]]: return [] +def suggest_favorite_query_with_template(text: str, arg: str) -> list[dict[str, Any]]: + favorite_arguments = arg.split(maxsplit=1) + if not favorite_arguments or (len(favorite_arguments) == 1 and not text[-1].isspace()): + return [{'type': 'favoritequery'}] + + name = favorite_arguments[0] + argument_text = favorite_arguments[1] if len(favorite_arguments) == 2 else '' + try: + arguments = shlex.split(argument_text) + except ValueError: + return [] + + used_keys: set[str] = set() + trailing_space = text[-1].isspace() + index = 0 + while index < len(arguments): + argument = arguments[index] + is_current = index == len(arguments) - 1 and not trailing_space + + if argument == '--': + return [] + if argument.startswith('--'): + option = argument[2:] + if '=' in option: + key, _value = option.split('=', 1) + used_keys.add(key) + if is_current: + return [] + elif is_current: + break + elif index + 1 >= len(arguments) or arguments[index + 1].startswith('--'): + return [] + else: + used_keys.add(option) + index += 1 + if index == len(arguments) - 1 and not trailing_space: + return [] + elif is_current: + if argument.startswith('-'): + break + return [] + index += 1 + + return [{'type': 'favoritequery_template_key', 'name': name, 'used_keys': used_keys}] + + def suggest_based_on_last_token( token: str | Token | None, text_before_cursor: str, diff --git a/mycli/packages/special/favoritequeries.py b/mycli/packages/special/favoritequeries.py index acde20b67..130839313 100644 --- a/mycli/packages/special/favoritequeries.py +++ b/mycli/packages/special/favoritequeries.py @@ -1,5 +1,64 @@ from __future__ import annotations +import re + +from jinja2 import meta, nodes +from jinja2.sandbox import SandboxedEnvironment + +favorite_query_template_environment = SandboxedEnvironment(autoescape=False) +favorite_query_variable_pattern = re.compile(r'^[A-Za-z_][A-Za-z0-9_-]*$') + + +def analyze_favorite_query_template(query: str) -> tuple[set[str], bool]: + """Return statically referenced keys and whether ``kv`` is used dynamically.""" + parsed_template = favorite_query_template_environment.parse(query) + if 'kv' not in meta.find_undeclared_variables(parsed_template): + return set(), False + + keys: set[str] = set() + called_attributes: set[int] = set() + accessed_names: set[int] = set() + dynamic_access = False + for call in parsed_template.find_all(nodes.Call): + called = call.node + if isinstance(called, nodes.Getattr) and isinstance(called.node, nodes.Name) and called.node.name == 'kv': + called_attributes.add(id(called)) + accessed_names.add(id(called.node)) + if called.attr == 'get' and call.args: + key = call.args[0] + if isinstance(key, nodes.Const) and isinstance(key.value, str): + keys.add(key.value) + else: + dynamic_access = True + else: + dynamic_access = True + + for attribute in parsed_template.find_all(nodes.Getattr): + if isinstance(attribute.node, nodes.Name) and attribute.node.name == 'kv': + accessed_names.add(id(attribute.node)) + if id(attribute) not in called_attributes: + keys.add(attribute.attr) + + for item in parsed_template.find_all(nodes.Getitem): + if isinstance(item.node, nodes.Name) and item.node.name == 'kv': + accessed_names.add(id(item.node)) + key = item.arg + if isinstance(key, nodes.Const) and isinstance(key.value, str): + keys.add(key.value) + else: + dynamic_access = True + + if any(name.name == 'kv' and id(name) not in accessed_names for name in parsed_template.find_all(nodes.Name)): + dynamic_access = True + + return keys, dynamic_access + + +def find_favorite_query_template_keys(query: str) -> set[str]: + """Return statically referenced keys from the template's ``kv`` dictionary.""" + keys, _dynamic_access = analyze_favorite_query_template(query) + return keys + class FavoriteQueries: section_name: str = "favorite_queries" @@ -14,11 +73,12 @@ class FavoriteQueries: # List all favorite queries. > /f - ╒════════╤═══════════════════════════════════════╕ - │ Name │ Query │ - ╞════════╪═══════════════════════════════════════╡ - │ simple │ SELECT * FROM abc where a is not NULL │ - ╘════════╧═══════════════════════════════════════╛ + ╒═══════════╤══════════════════════════════════════════════════╕ + │ Name │ Query │ + ╞═══════════╪══════════════════════════════════════════════════╡ + │ simple │ SELECT * FROM abc where a is not NULL │ + │ find_user │ SELECT * FROM users WHERE name = '{{ kv.name }}' │ + ╘═══════════╧══════════════════════════════════════════════════╛ # Run a favorite query. > /f simple @@ -28,6 +88,18 @@ class FavoriteQueries: │ 日本語 │ 日本語 │ ╘════════╧════════╛ + # Run a favorite query containing {{ kv.name }} in the template: + > /f find_user --name=henry + > /f find_user --name henry + + # Run a favorite query containing $1 in the template: + > /f find_user henry + + # Use -- to disambiguate positional parameters, especially if + # the positional value starts with a dash. + > /f query --key=value -- positional-value + > /f query -- --positional-value-which-looks-like-a-flag-- + # Delete a favorite query. > /fd simple simple: Deleted. diff --git a/mycli/packages/special/iocommands.py b/mycli/packages/special/iocommands.py index e5ee06360..01ae3d7f2 100644 --- a/mycli/packages/special/iocommands.py +++ b/mycli/packages/special/iocommands.py @@ -8,9 +8,11 @@ import subprocess from time import sleep from typing import Any, Generator +from uuid import uuid4 import click from configobj import ConfigObj +from jinja2 import TemplateError from prompt_toolkit.formatted_text import ANSI, FormattedText, to_plain_text from pymysql.cursors import Cursor import pyperclip @@ -20,7 +22,12 @@ from mycli.packages.interactive_utils import confirm_destructive_query from mycli.packages.special.delimitercommand import DelimiterCommand from mycli.packages.special.dsn_aliases import INVALID_DSN_ALIAS_ERROR, DsnAliases, is_valid_dsn_alias -from mycli.packages.special.favoritequeries import FavoriteQueries +from mycli.packages.special.favoritequeries import ( + FavoriteQueries, + analyze_favorite_query_template, + favorite_query_template_environment, + favorite_query_variable_pattern, +) from mycli.packages.special.main import COMMANDS as SPECIAL_COMMANDS from mycli.packages.special.main import ArgType, SpecialCommandAlias, special_command from mycli.packages.special.main import execute as special_execute @@ -52,6 +59,10 @@ SHOW_WARNINGS_ENABLED: bool = False +class FavoriteQueryArgumentError(ValueError): + pass + + def set_favorite_queries(config): global favoritequeries favoritequeries = FavoriteQueries(config) @@ -339,7 +350,7 @@ def set_redirect(command_part: str | None, file_operator_part: str | None, file_ @special_command( "\\f", - "/f [name [args..]]", + "/f [name [args..] [--key=value]]", "List or execute favorite queries.", arg_type=ArgType.PARSED_QUERY, case_sensitive=True, @@ -351,17 +362,33 @@ def execute_favorite_query(cur: Cursor, arg: str, **_) -> Generator[SQLResult, N # Parse out favorite name and optional substitution parameters name, _separator, arg_str = arg.partition(" ") - args = shlex.split(arg_str) + try: + args, template_values = parse_favorite_query_args(arg_str) + except ValueError as exc: + yield SQLResult(status=f'Invalid favorite query arguments: {exc}') + return query = FavoriteQueries.instance.get(name) if query is None: message = f"No favorite query: {name}" yield SQLResult(status=message) else: - query, arg_error = subst_favorite_query_args(query, args) + query, positional_values, arg_error = prepare_favorite_query_args(query, args) if query is None: yield SQLResult(status=arg_error) else: + try: + query = render_favorite_query(query, template_values) + except TemplateError as exc: + yield SQLResult(status=f'Favorite query template error: {exc}') + return + except FavoriteQueryArgumentError as exc: + yield SQLResult(status=f'Invalid favorite query arguments: {exc}') + return + except Exception as exc: + yield SQLResult(status=f'Favorite query template error: {exc}') + return + query = restore_favorite_query_args(query, positional_values) for sql in sqlparse.split(query): sql = sql.rstrip(";") preamble = f"> {sql}" if is_show_favorite_query() else None @@ -384,6 +411,49 @@ def execute_favorite_query(cur: Cursor, arg: str, **_) -> Generator[SQLResult, N yield SQLResult(preamble=preamble) +def parse_favorite_query_args(arg_str: str) -> tuple[list[str], dict[str, str]]: + """Split favorite query arguments into positional and template values.""" + tokens = shlex.split(arg_str) + positional: list[str] = [] + template_values: dict[str, str] = {} + parse_options = True + index = 0 + + while index < len(tokens): + token = tokens[index] + if parse_options and token == '--': + parse_options = False + elif parse_options and token.startswith('--'): + option = token[2:] + if '=' in option: + name, value = option.split('=', 1) + else: + name = option + index += 1 + if index >= len(tokens) or tokens[index].startswith('--'): + raise ValueError(f'option --{name} requires a value') + value = tokens[index] + + if not favorite_query_variable_pattern.fullmatch(name): + raise ValueError(f'invalid template variable name: {name or token}') + if name in template_values: + raise ValueError(f'duplicate template variable: {name}') + template_values[name] = value + else: + positional.append(token) + index += 1 + + return positional, template_values + + +def render_favorite_query(query: str, template_values: dict[str, str]) -> str: + referenced_keys, dynamic_access = analyze_favorite_query_template(query) + unused_variables = sorted(set(template_values) - referenced_keys) + if unused_variables and not dynamic_access: + raise FavoriteQueryArgumentError(f'unused template variable: {", ".join(unused_variables)}') + return favorite_query_template_environment.from_string(query).render(kv=template_values) + + def list_favorite_queries() -> list[SQLResult]: """List of all favorite queries.""" @@ -397,20 +467,44 @@ def list_favorite_queries() -> list[SQLResult]: return [SQLResult(header=header, rows=rows, status=status)] -def subst_favorite_query_args(query: str, args: list[str]) -> list[str | None]: - """replace positional parameters ($1...$N) in query.""" +def restore_favorite_query_args(query: str, positional_values: dict[str, str]) -> str: + for marker, value in positional_values.items(): + query = query.replace(marker, value) + return query + + +def prepare_favorite_query_args(query: str, args: list[str]) -> tuple[str | None, dict[str, str], str | None]: + """Replace positional parameters with markers that survive template rendering.""" + marker_prefix = f'__mycli_favorite_arg_{uuid4().hex}_' + + positional_values: dict[str, str] = {} for idx, val in enumerate(args): subst_var = "$" + str(idx + 1) if subst_var not in query: - return [None, "query does not have substitution parameter " + subst_var + ":\n " + query] + display_query = restore_favorite_query_args(query, positional_values) + error = "query does not have substitution parameter " + subst_var + ":\n " + display_query + return (None, {}, error) - query = query.replace(subst_var, val) + marker = f'{marker_prefix}{idx + 1}__' + query = query.replace(subst_var, marker) + positional_values[marker] = val match = re.search(r"\$\d+", query) if match: - return [None, "missing substitution for " + match.group(0) + " in query:\n " + query] + display_query = restore_favorite_query_args(query, positional_values) + error = "missing substitution for " + match.group(0) + " in query:\n " + display_query + return (None, {}, error) + + return (query, positional_values, None) + + +def subst_favorite_query_args(query: str, args: list[str]) -> list[str | None]: + """Replace positional parameters ($1...$N) in query.""" + prepared_query, positional_values, error = prepare_favorite_query_args(query, args) + if prepared_query is None: + return [None, error] - return [query, None] + return [restore_favorite_query_args(prepared_query, positional_values), None] @special_command( diff --git a/mycli/sqlcompleter.py b/mycli/sqlcompleter.py index 362ea733f..f085f4040 100644 --- a/mycli/sqlcompleter.py +++ b/mycli/sqlcompleter.py @@ -6,6 +6,7 @@ import re from typing import Any, Collection, Generator, Iterable, Literal +from jinja2 import TemplateError from prompt_toolkit.completion import CompleteEvent, Completer, Completion from prompt_toolkit.completion.base import Document from pygments.lexers._mysql_builtins import MYSQL_DATATYPES, MYSQL_FUNCTIONS, MYSQL_KEYWORDS @@ -15,7 +16,11 @@ from mycli.packages.filepaths import complete_path, parse_path, suggest_path from mycli.packages.special import llm from mycli.packages.special.dsn_aliases import DsnAliases -from mycli.packages.special.favoritequeries import FavoriteQueries +from mycli.packages.special.favoritequeries import ( + FavoriteQueries, + favorite_query_variable_pattern, + find_favorite_query_template_keys, +) from mycli.packages.special.main import COMMANDS as SPECIAL_COMMANDS from mycli.packages.sql_utils import extract_columns_from_select, extract_tables, last_word @@ -1680,6 +1685,27 @@ def get_completions( ) completions.extend([(*x, rank) for x in queries_m]) + elif suggestion['type'] == 'favoritequery_template_key': + if hasattr(FavoriteQueries, 'instance') and hasattr(FavoriteQueries.instance, 'get'): + query = FavoriteQueries.instance.get(suggestion['name']) + if query is not None: + try: + template_keys = find_favorite_query_template_keys(query) + except TemplateError: + continue + used_keys = suggestion['used_keys'] + candidates = [ + f'--{key}=' for key in sorted(template_keys - used_keys) if favorite_query_variable_pattern.fullmatch(key) + ] + keys_m = self.find_matches( + word_before_cursor, + candidates, + start_only=True, + fuzzy=False, + text_before_cursor=document.text_before_cursor, + ) + completions.extend([(*x, rank) for x in keys_m]) + elif suggestion["type"] == "table_format": formats_m = self.find_matches( word_before_cursor, diff --git a/pyproject.toml b/pyproject.toml index e0bf8807e..d045c0c22 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -42,6 +42,7 @@ dependencies = [ "pyfzf ~= 0.3.1", "rapidfuzz ~= 3.14.3", "keyring ~= 25.7.0", + "jinja2 ~= 3.1.6", "yaspin ~= 3.4.0", ] diff --git a/test/features/fixture_data/help_commands.txt b/test/features/fixture_data/help_commands.txt index de526daaa..2dfc433bd 100644 --- a/test/features/fixture_data/help_commands.txt +++ b/test/features/fixture_data/help_commands.txt @@ -9,7 +9,7 @@ | /dt | | /dt[+] [table] | List or describe tables. | | /edit | /e | /edit | \edit | Edit query with editor (uses $VISUAL or $EDITOR). | | /exit | /q | /exit | Exit. | -| /f | | /f [name [args..]] | List or execute favorite queries. | +| /f | | /f [name [args..] [--key=value]] | List or execute favorite queries. | | /fd | | /fd | Delete a favorite query. | | /fs | | /fs | Save a favorite query. | | \g | | \g | Display query results (mnemonic: go). | diff --git a/test/features/named_queries.feature b/test/features/named_queries.feature index 5e681ec4c..28ca581e8 100644 --- a/test/features/named_queries.feature +++ b/test/features/named_queries.feature @@ -22,3 +22,13 @@ Feature: named queries: then we see the named query with parameters fail with missing parameters when we use named query with too many parameters then we see the named query with parameters fail with extra parameters + + Scenario: use a named query as a Jinja template + When we connect to test database + then we see database connected + when we save a templated named query + then we see the named query saved + when we use a templated named query with attached values + then we see the attached template values rendered + when we use a templated named query with split values + then we see the split template values rendered diff --git a/test/features/steps/named_queries.py b/test/features/steps/named_queries.py index ea53234cb..246d22fbd 100644 --- a/test/features/steps/named_queries.py +++ b/test/features/steps/named_queries.py @@ -87,3 +87,28 @@ def step_use_named_query_with_too_many_parameters(context): def step_see_named_query_with_parameters_fail_with_extra_parameters(context): """Wait to see select output.""" wrappers.expect_exact(context, "query does not have substitution parameter $4:", timeout=2) + + +@when("we save a templated named query") +def step_save_templated_named_query(context): + context.cli.sendline("\\fs template SELECT '{{ kv.user }}', '$1'") + + +@when("we use a templated named query with attached values") +def step_use_templated_named_query_with_attached_values(context): + context.cli.sendline("\\f template positional --user=henry") + + +@then("we see the attached template values rendered") +def step_see_attached_template_values_rendered(context): + wrappers.expect_exact(context, "SELECT 'henry', 'positional'", timeout=2) + + +@when("we use a templated named query with split values") +def step_use_templated_named_query_with_split_values(context): + context.cli.sendline('\\f template second --user "Henry Ford"') + + +@then("we see the split template values rendered") +def step_see_split_template_values_rendered(context): + wrappers.expect_exact(context, "SELECT 'Henry Ford', 'second'", timeout=2) diff --git a/test/myclirc b/test/myclirc index 056327778..e36c38f37 100644 --- a/test/myclirc +++ b/test/myclirc @@ -495,8 +495,32 @@ matching-bracket.other = '#000000 bg:#aacccc' # Favorite queries. # You can add your favorite queries here. They will be available in the # REPL when you type `/f` or `/f `. +# Define multiine favorite queries using '''triplequotes'''. [favorite_queries] # example = "SELECT * FROM example_table WHERE id = 1" +# +# Favorite queries can use the Jinja2 templating language. Named argument +# values are available as keys in a dictionary named "kv". +# Run this example templated query with: /f find_user --name=henry +# Any keys refereed to by name will be available as completions. +# find_user = "SELECT * FROM users WHERE name = '{{ kv.name }}'" +# +# Jinja2 parameters are optional. +# Run this example templated query with: /f recent_user --name=henry +# or alternatively, with different behavior: /f recent_user +# recent_user = "SELECT * FROM users WHERE create_date >= NOW() - INTERVAL 1 DAY {% if kv.name %} AND name = '{{ kv.name }}' {% endif %}" +# +# Since "kv" is a Python dictionary, collisions with existing dictionary +# methods will require extra syntax, as will keys with internal dashes: +# example = "SELECT * FROM example_table WHERE id = {{ kv['get'] }}" +# example = "SELECT * FROM example_table WHERE id = {{ kv['my-id'] }}" +# +# Mandatory positional arguments are also supported: +# Run this example templated query with: /f find_user henry +# find_user = "SELECT * FROM users WHERE name = '$1'" +# +# Plase ambiguous positional arguments after "--": +# /f find_user -- --unusual-name-- # Initial commands to execute when connecting to any database. [init-commands] diff --git a/test/pytests/test_completion_engine.py b/test/pytests/test_completion_engine.py index cafdab991..d890fd8c6 100644 --- a/test/pytests/test_completion_engine.py +++ b/test/pytests/test_completion_engine.py @@ -859,6 +859,29 @@ def test_suggest_type_handles_parser_results_shorter_than_cursor(monkeypatch): ('\\f ', [{'type': 'favoritequery'}]), ('\\fs ', [{'type': 'favoritequery'}]), ('\\fd ', [{'type': 'favoritequery'}]), + ('/f report', [{'type': 'favoritequery'}]), + ('/f report ', [{'type': 'favoritequery_template_key', 'name': 'report', 'used_keys': set()}]), + ('\\f report --u', [{'type': 'favoritequery_template_key', 'name': 'report', 'used_keys': set()}]), + ('/f report -', [{'type': 'favoritequery_template_key', 'name': 'report', 'used_keys': set()}]), + ('/f report positional', []), + ('/f report positional ', [{'type': 'favoritequery_template_key', 'name': 'report', 'used_keys': set()}]), + ('/f report --user=', []), + ( + '/f report --user=henry ', + [{'type': 'favoritequery_template_key', 'name': 'report', 'used_keys': {'user'}}], + ), + ('/f report --user ', []), + ('/f report --user henry', []), + ( + '/f report --user henry ', + [{'type': 'favoritequery_template_key', 'name': 'report', 'used_keys': {'user'}}], + ), + ( + '/f report --start-date 2026-08-01 ', + [{'type': 'favoritequery_template_key', 'name': 'report', 'used_keys': {'start-date'}}], + ), + ('/f report -- --user', []), + ('/f report --user="henry', []), ('\\dt ', [{'type': 'table', 'schema': []}, {'type': 'view', 'schema': []}, {'type': 'schema'}]), ('\\dt+ ', [{'type': 'table', 'schema': []}, {'type': 'view', 'schema': []}, {'type': 'schema'}]), ('\\. ', [{'type': 'file_name'}]), diff --git a/test/pytests/test_special_iocommands.py b/test/pytests/test_special_iocommands.py index 14f983386..4c036a05e 100644 --- a/test/pytests/test_special_iocommands.py +++ b/test/pytests/test_special_iocommands.py @@ -4,6 +4,7 @@ import os from pathlib import Path import platform +import re import stat import subprocess import tempfile @@ -12,11 +13,13 @@ from typing import Any, Generator from unittest.mock import patch +from jinja2 import TemplateError from pymysql import ProgrammingError import pytest import mycli.packages.special from mycli.packages.special import iocommands +from mycli.packages.special.favoritequeries import analyze_favorite_query_template, find_favorite_query_template_keys from mycli.packages.sqlresult import SQLResult from test.utils import TEMPFILE_PREFIX, db_connection, dbtest, send_ctrl_c @@ -675,6 +678,183 @@ def test_execute_favorite_query_returns_header_for_result_sets(monkeypatch) -> N assert results[0].rows is cursor +@pytest.mark.parametrize( + ('arg_str', 'expected_positional', 'expected_template_values'), + ( + ('', [], {}), + ('first second', ['first', 'second'], {}), + ('--user=henry', [], {'user': 'henry'}), + ('--user "Henry Ford"', [], {'user': 'Henry Ford'}), + ('--empty=', [], {'empty': ''}), + ('--start-date 2026-08-01', [], {'start-date': '2026-08-01'}), + ( + '--start-date=one --start_date=two', + [], + {'start-date': 'one', 'start_date': 'two'}, + ), + ( + 'first --user=henry second -- --literal --other=value', + ['first', 'second', '--literal', '--other=value'], + {'user': 'henry'}, + ), + ), +) +def test_parse_favorite_query_args( + arg_str: str, + expected_positional: list[str], + expected_template_values: dict[str, str], +) -> None: + assert iocommands.parse_favorite_query_args(arg_str) == (expected_positional, expected_template_values) + + +@pytest.mark.parametrize( + ('arg_str', 'message'), + ( + ('--user', 'option --user requires a value'), + ('--user --other=value', 'option --user requires a value'), + ('--=value', 'invalid template variable name: --=value'), + ('--1name=value', 'invalid template variable name: 1name'), + ('--start-date=one --start-date=two', 'duplicate template variable: start-date'), + ), +) +def test_parse_favorite_query_args_rejects_invalid_options(arg_str: str, message: str) -> None: + with pytest.raises(ValueError, match=re.escape(message)): + iocommands.parse_favorite_query_args(arg_str) + + +def test_parse_favorite_query_args_rejects_malformed_quoting() -> None: + with pytest.raises(ValueError, match='No closing quotation'): + iocommands.parse_favorite_query_args('--user="henry') + + +def test_render_favorite_query_supports_jinja_features_and_missing_values() -> None: + query = '{% for item in kv.item_list.split(",") %}{{ item|upper }} {% endfor %}{% if kv.enabled %}enabled{% endif %} {{ kv.missing }}' + + assert iocommands.render_favorite_query(query, {'item_list': 'one,two', 'enabled': 'yes'}) == 'ONE TWO enabled ' + + +def test_find_favorite_query_template_keys_excludes_jinja_locals_and_globals() -> None: + query = '{% set local = kv.user %}{% for item in range(kv.limit|int) %}{{ local }} {{ item }}{% endfor %}' + + assert find_favorite_query_template_keys(query) == {'user', 'limit'} + + +def test_find_favorite_query_template_keys_supports_dictionary_access_forms() -> None: + query = "{{ kv.range }} {{ kv['dict'] }} {{ kv.get('namespace') }} {{ range(2) }} {{ kv.items() }}" + + assert find_favorite_query_template_keys(query) == {'range', 'dict', 'namespace'} + assert ( + iocommands.render_favorite_query( + "{{ kv.range }} {{ kv['dict'] }} {{ kv.get('namespace') }}", + {'range': 'one', 'dict': 'two', 'namespace': 'three'}, + ) + == 'one two three' + ) + + +@pytest.mark.parametrize( + ('query', 'expected_keys'), + ( + ('{{ kv }}', set()), + ('{% for key, value in kv.items() %}{{ key }}={{ value }} {% endfor %}', set()), + ('{% set values = kv %}{{ values.user }}', set()), + ('{{ kv[kv.which] }}', {'which'}), + ('{{ kv.get(kv.which) }}', {'which'}), + ), +) +def test_analyze_favorite_query_template_detects_dynamic_access(query: str, expected_keys: set[str]) -> None: + assert analyze_favorite_query_template(query) == (expected_keys, True) + + +def test_render_favorite_query_rejects_unused_values() -> None: + with pytest.raises(iocommands.FavoriteQueryArgumentError, match='unused template variable: extra, unused'): + iocommands.render_favorite_query('select {{ kv.used }}', {'used': '1', 'unused': '2', 'extra': '3'}) + + +@pytest.mark.parametrize( + ('query', 'template_values', 'expected'), + ( + ('{{ kv.user }} {{ kv }}', {'user': 'henry', 'extra': 'value'}, "henry {'user': 'henry', 'extra': 'value'}"), + ( + '{% for key, value in kv.items()|sort %}{{ key }}={{ value }} {% endfor %}', + {'user': 'henry', 'role': 'admin'}, + 'role=admin user=henry ', + ), + ('{% set values = kv %}{{ values.user }}', {'user': 'henry'}, 'henry'), + ('{{ kv[kv.which] }}', {'which': 'user', 'user': 'henry'}, 'henry'), + ('{{ kv.get(kv.which) }}', {'which': 'user', 'user': 'henry'}, 'henry'), + ), +) +def test_render_favorite_query_allows_dynamic_template_values( + query: str, + template_values: dict[str, str], + expected: str, +) -> None: + assert iocommands.render_favorite_query(query, template_values) == expected + + +def test_render_favorite_query_uses_sandbox() -> None: + with pytest.raises(TemplateError, match='unsafe'): + iocommands.render_favorite_query("{{ ''.__class__.__mro__ }}", {}) + + +def test_execute_favorite_query_renders_named_and_positional_values(monkeypatch) -> None: + query = """select '$1', '{{ kv.user }}', '{{ kv["start-date"] }}', '{{ kv.literal }}'""" + monkeypatch.setattr(iocommands.FavoriteQueries, 'instance', FakeFavoriteQueries({'report': query}), raising=False) + cursor = FakeCursor() + + results = list( + iocommands.execute_favorite_query( + cursor, + "report positional --user=henry --start-date 2026-08-01 --literal='$1'", + ) + ) + + expected_query = "select 'positional', 'henry', '2026-08-01', '$1'" + assert cursor.executed == [expected_query] + assert results[0].preamble == f'> {expected_query}' + + +def test_execute_favorite_query_does_not_render_runtime_values_as_jinja(monkeypatch) -> None: + query = "select '$1', '{{ kv.literal }}'" + monkeypatch.setattr(iocommands.FavoriteQueries, 'instance', FakeFavoriteQueries({'report': query}), raising=False) + cursor = FakeCursor() + + results = list(iocommands.execute_favorite_query(cursor, "report '{{ 2 * 3 }}' --literal='$1'")) + + expected_query = "select '{{ 2 * 3 }}', '$1'" + assert cursor.executed == [expected_query] + assert results[0].preamble == f'> {expected_query}' + + +@pytest.mark.parametrize( + ('query', 'arg', 'status_prefix'), + ( + ('select {{ kv.user }}', 'report --user', 'Invalid favorite query arguments:'), + ('select {{ kv.user }}', 'report --unused=value', 'Invalid favorite query arguments:'), + ('select {{', 'report', 'Favorite query template error:'), + ("select {{ ''.__class__.__mro__ }}", 'report', 'Favorite query template error:'), + ('select {{ range(10**100) }}', 'report', 'Favorite query template error:'), + ('select {{ range(1, 2, 0) }}', 'report', 'Favorite query template error:'), + ('select {{ kv.user + 1 }}', 'report --user=2', 'Favorite query template error:'), + ), +) +def test_execute_favorite_query_reports_template_argument_errors_without_execution( + monkeypatch: pytest.MonkeyPatch, + query: str, + arg: str, + status_prefix: str, +) -> None: + monkeypatch.setattr(iocommands.FavoriteQueries, 'instance', FakeFavoriteQueries({'report': query}), raising=False) + cursor = FakeCursor() + + results = list(iocommands.execute_favorite_query(cursor, arg)) + + assert results[0].status is not None + assert results[0].status.startswith(status_prefix) + assert cursor.executed == [] + + def test_list_substitute_save_delete_and_redirect_state(tmp_path: Path, monkeypatch) -> None: empty_favorites = FakeFavoriteQueries() monkeypatch.setattr(iocommands.FavoriteQueries, 'instance', empty_favorites, raising=False) @@ -690,6 +870,7 @@ def test_list_substitute_save_delete_and_redirect_state(tmp_path: Path, monkeypa assert rows_result.status == '' assert iocommands.subst_favorite_query_args('select $1', ['x']) == ['select x', None] + assert iocommands.subst_favorite_query_args('select $1, $2', ['$2', 'second']) == ['select $2, second', None] assert iocommands.subst_favorite_query_args('select 1', ['x']) == [None, 'query does not have substitution parameter $1:\n select 1'] assert iocommands.subst_favorite_query_args('select $1, $2', ['x']) == [None, 'missing substitution for $2 in query:\n select x, $2'] diff --git a/test/pytests/test_sqlcompleter.py b/test/pytests/test_sqlcompleter.py index f3ca1be55..b67c7277b 100644 --- a/test/pytests/test_sqlcompleter.py +++ b/test/pytests/test_sqlcompleter.py @@ -511,6 +511,50 @@ def test_get_completions_branch_specific_suggestions(monkeypatch, suggestion, se assert expected in result +def test_get_completions_favorite_query_template_keys(monkeypatch) -> None: + queries = { + 'report': "select {{ kv.user }}, {{ kv.start_date }}, {{ kv['start-date'] }}, {{ range(2) }}, {{ kv.range }}", + } + monkeypatch.setattr( + mycli.sqlcompleter.FavoriteQueries, + 'instance', + SimpleNamespace(list=lambda: list(queries), get=queries.get), + raising=False, + ) + completer = make_completer() + + blank_text = '/f report ' + blank = list(completer.get_completions(Document(text=blank_text, cursor_position=len(blank_text)), None)) + partial_text = '/f report --u' + partial = list(completer.get_completions(Document(text=partial_text, cursor_position=len(partial_text)), None)) + dashed_text = '/f report --start-' + dashed = list(completer.get_completions(Document(text=dashed_text, cursor_position=len(dashed_text)), None)) + used_text = '/f report --user=henry ' + used = list(completer.get_completions(Document(text=used_text, cursor_position=len(used_text)), None)) + used_dashed_text = '/f report --start-date=2026-08-03 ' + used_dashed = list(completer.get_completions(Document(text=used_dashed_text, cursor_position=len(used_dashed_text)), None)) + + assert [completion.text for completion in blank] == ['--range=', '--start-date=', '--start_date=', '--user='] + assert [(completion.text, completion.start_position) for completion in partial] == [('--user=', -3)] + assert [(completion.text, completion.start_position) for completion in dashed] == [('--start-date=', -8)] + assert [completion.text for completion in used] == ['--range=', '--start-date=', '--start_date='] + assert [completion.text for completion in used_dashed] == ['--range=', '--start_date=', '--user='] + + +@pytest.mark.parametrize('query', [None, '{{ invalid']) +def test_get_completions_favorite_query_template_keys_fail_quietly(monkeypatch, query) -> None: + monkeypatch.setattr( + mycli.sqlcompleter.FavoriteQueries, + 'instance', + SimpleNamespace(list=lambda: ['report'], get=lambda name: query), + raising=False, + ) + completer = make_completer() + text = '/f report ' + + assert list(completer.get_completions(Document(text=text, cursor_position=len(text)), None)) == [] + + def test_get_completions_llm_branch_with_and_without_current_word(monkeypatch) -> None: tokens_seen: list[list[str]] = []