Let favorite queries use Jinja2 templates

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`.

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 <query_name>`.
+# 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]

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,

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.