tinybird-cli 6.5.6.dev0__py3-none-any.whl → 6.5.8.dev0__py3-none-any.whl

tinybird/sql_template.py

@@ -1,5 +1,4 @@
 import ast
-import builtins
 import linecache
 import logging
 import re
@@ -18,7 +17,7 @@ from tinybird.context import (
 )
 
 from .datatypes import testers
-from .tornado_template import VALID_CUSTOM_FUNCTION_NAMES, SecurityException, Template
+from .tornado_template import SAFE_BUILTINS, VALID_CUSTOM_FUNCTION_NAMES, SecurityException, Template
 
 VALID_ACTIVATE_FEATURES = frozenset(["analyzer", "parallel_replicas", "optimize_aggregation_in_order"])
 
@@ -84,6 +83,7 @@ class SQLTemplateException(ValueError):
 JOB_TIMESTAMP_PARAM = "job_timestamp"
 DEFAULT_PARAM_NAMES = ["format", "q"]
 RESERVED_PARAM_NAMES = [
+    "__tb__csv_header",
     "__tb__semver",
     "debug_source_tables",
     "debug",
@@ -142,8 +142,7 @@ def transform_type(
             if x is None:
                 if defined:
                     raise SQLTemplateException(REQUIRED_PARAM_NOT_DEFINED, documentation="/cli/advanced-templates.html")
-                else:
-                    return None
+                return None
         if tester == "String":
             if x is not None:
                 return transform(x)
@@ -167,12 +166,11 @@ def _and(*args, **kwargs):
         tk = k.rsplit("__", 1)
         if len(tk) == 1:
             return "="
-        else:
-            if tk[1] in operands:
-                return operands[tk[1]]
-            raise SQLTemplateException(
-                f"operand {tk[1]} not supported", documentation="/cli/advanced-templates.html#sql_and"
-            )
+        if tk[1] in operands:
+            return operands[tk[1]]
+        raise SQLTemplateException(
+            f"operand {tk[1]} not supported", documentation="/cli/advanced-templates.html#sql_and"
+        )
 
     return Expression(
         " and ".join([f"{_name(k)} {_op(k)} {expression_wrapper(v, k)}" for k, v in kwargs.items() if v is not None])
@@ -277,8 +275,7 @@ def columns(x, default=None, fn=None):
 
     if fn:
         return Expression(",".join(f"{fn}({str(column(c, c))}) as {c}" for c in _columns))
-    else:
-        return Expression(",".join(str(column(c, c)) for c in _columns))
+    return Expression(",".join(str(column(c, c)) for c in _columns))
 
 
 def column(x, default=None):
@@ -368,11 +365,10 @@ def boolean(x, default=None):
         if default is None:
             return 0
         return boolean(default)
-    elif isinstance(x, Placeholder):
+    if isinstance(x, Placeholder):
         return boolean(default)
-    elif isinstance(x, str):
-        if x == "0" or x.lower() == "false":
-            return 0
+    if isinstance(x, str) and (x == "0" or x.lower() == "false"):
+        return 0
 
     return int(bool(x))
 
@@ -403,8 +399,7 @@ def array_type(types):
                         raise SQLTemplateException(
                             REQUIRED_PARAM_NOT_DEFINED, documentation="/cli/advanced-templates.html"
                         )
-                    else:
-                        return None
+                    return None
             values = []
             list_values = x if type(x) == list else x.split(",")  # noqa: E721
             for i, t in enumerate(list_values):
@@ -428,22 +423,29 @@ def array_type(types):
     return _f
 
 
-def sql_unescape(x, what=""):
+def sql_unescape(x: str, what: str = "") -> Expression:
     """
     unescapes specific characters in a string. It allows to allow some
     special characters to be used, for example in like condictionals
 
     {{sql_unescape(String(like_filter), '%')}}
 
+    `what` must be given and must not contain quotes or backslashes: unescaping
+    those would let the caller close the string literal, so they stay escaped.
 
     >>> sql_unescape('testing%', '%')
     "'testing%'"
     >>> sql_unescape('testing%', '$')
     "'testing\\\\%'"
     >>> sql_unescape('testing"')
-    '\\'testing"\\''
+    '\\'testing\\\\"\\''
+    >>> sql_unescape("x' OR 1=1", "'")
+    "'x\\\\' OR 1=1'"
     """
-    return Expression("'" + sqlescape_for_string_expression(x).replace(f"\\{what}", what) + "'")
+    escaped = sqlescape_for_string_expression(x)
+    if what and not any(c in what for c in "'\"\\"):
+        escaped = escaped.replace(f"\\{what}", what)
+    return Expression("'" + escaped + "'")
 
 
 def split_to_array(x, default="", separator: str = ","):
@@ -744,12 +746,11 @@ def __date_diff(
 
         if unit == "days":
             return int(diff / 86400)
-        elif unit == "hours":
+        if unit == "hours":
             return int(diff / 3600)
-        elif unit == "minutes":
+        if unit == "minutes":
             return int(diff / 60)
-        else:
-            return int(diff)
+        return int(diff)
     except Exception:
         if none_if_error:
             return None
@@ -1341,16 +1342,18 @@ sqlescape_for_string_expression = sqlescape_generator(
 )
 
 
-def escape_single_quote_str(s):
-    return "'" + s.replace("'", "''") + "'"
+def escape_single_quote_str(s: str) -> str:
+    # ClickHouse string literals treat backslash as an escape char, so it must be doubled too:
+    # otherwise a trailing backslash escapes the closing quote and the following input becomes SQL
+    return "'" + s.replace("\\", "\\\\").replace("'", "''") + "'"
 
 
 def expression_wrapper(x, name, escape_arrays: bool = False):
     if type(x) in (unicode_type, bytes, str):
         return "'" + sqlescape_for_string_expression(x) + "'"
-    elif isinstance(x, Placeholder):
+    if isinstance(x, Placeholder):
         return "'__no_value__'"
-    elif isinstance(x, Comment):
+    if isinstance(x, Comment):
         return "-- {x} \n"
     if x is None:
         truncated_name = name[:20] + "..." if len(name) > 20 else name
@@ -1361,10 +1364,7 @@ def expression_wrapper(x, name, escape_arrays: bool = False):
         logging.warning(f"expression_wrapper -> list :{x}:")
 
         try:
-            result = (
-                f"[{','.join(escape_single_quote_str(item) if isinstance(item, str) else str(item) for item in x)}]"
-            )
-            return result
+            return f"[{','.join(escape_single_quote_str(item) if isinstance(item, str) else str(item) for item in x)}]"
         except Exception as e:
             logging.error(f"Error escaping array: {e}")
     return x
@@ -1383,7 +1383,7 @@ _namespace = {
     # they raise a pretty non understandable error but if someone
     # is using them they know what they are trying to do
     # read https://anee.me/escaping-python-jails-849c65cf306e on how to escape from python jails
-    "__buildins__": {},
+    "__builtins__": SAFE_BUILTINS,
     "__import__": {},
     "__debug__": {},
     "__doc__": {},
@@ -1396,10 +1396,13 @@ _namespace = {
 }
 
 
-reserved_vars = {"_tt_tmp", "_tt_append", "isinstance", "str", "error", "custom_error", *list(vars(builtins))}
+# Built from the builtins templates can actually resolve: names missing from SAFE_BUILTINS (`type`, `dir`,
+# `property`, `object`, ...) can only be template parameters, so they must be discovered as such to get a
+# value (None or a test-mode placeholder) instead of raising NameError.
+reserved_vars = {"_tt_tmp", "_tt_append", "isinstance", "str", "error", "custom_error", *SAFE_BUILTINS}
 for p in DEFAULT_PARAM_NAMES:  # we handle these in an specific manner
     reserved_vars.discard(p)  # `format` is part of builtins
-# Allow 'id' to be used as a template parameter - https://gitlab.com/tinybird/analytics/-/issues/19119
+# Allow 'id' as a template parameter name despite being a builtin - https://gitlab.com/tinybird/analytics/-/issues/19119
 reserved_vars.discard("id")
 error_vars = ["error", "custom_error"]
 
@@ -1425,20 +1428,18 @@ def generate(self, **kwargs) -> Tuple[str, TemplateExecutionResults]:
                 # secret available: Always use workspace secret regardless of test mode
                 template_execution_results.add_ch_param(x)
                 return Symbol("{" + sqlescape(x) + ": String}")
-            else:
-                # secret not available: Check test mode and defaults
-                is_test_mode = TB_SECRET_IN_TEST_MODE in template_execution_results
-                if default is not None:
-                    # Use provided default value
-                    return default
-                elif is_test_mode:
-                    # In test mode without default - return placeholder
-                    return Symbol("{" + sqlescape(x) + ": String}")
-                else:
-                    # Not in test mode, no secret, no default - raise error
-                    raise SQLTemplateException(
-                        f"Cannot access secret '{x}'. Check the secret exists in the Workspace and the token has the required scope."
-                    )
+            # secret not available: Check test mode and defaults
+            is_test_mode = TB_SECRET_IN_TEST_MODE in template_execution_results
+            if default is not None:
+                # Use provided default value
+                return default
+            if is_test_mode:
+                # In test mode without default - return placeholder
+                return Symbol("{" + sqlescape(x) + ": String}")
+            # Not in test mode, no secret, no default - raise error
+            raise SQLTemplateException(
+                f"Cannot access secret '{x}'. Check the secret exists in the Workspace and the token has the required scope."
+            )
         except Exception:
             raise SQLTemplateException(
                 f"Cannot access secret '{x}'. Check the secret exists in the Workspace and the token has the required scope."
@@ -1478,8 +1479,7 @@ def generate(self, **kwargs) -> Tuple[str, TemplateExecutionResults]:
         return Expression(f"-- disable {feature}\n")
 
     namespace.update(_namespace)
-    # This is to fix the issue https://gitlab.com/tinybird/analytics/-/issues/19119
-    # We need to do the override here because if we modify the _namespace, we would filter out parameters from the users
+    # Default it here, not in _namespace, so it stays visible to variable extraction.
     namespace.update({"id": None})
     namespace.update(kwargs)
     namespace.update(
@@ -1488,6 +1488,8 @@ def generate(self, **kwargs) -> Tuple[str, TemplateExecutionResults]:
             # the generated source code.
             "__name__": self.name.replace(".", "_"),
             "__loader__": ObjectDict(get_source=lambda name: self.code),
+            # Set again after kwargs so request parameters can't override it
+            "__builtins__": SAFE_BUILTINS,
             "max_threads": set_max_threads,
             "tb_secret": set_tb_secret,
             "tb_var": set_tb_secret,
@@ -1522,14 +1524,12 @@ def generate(self, **kwargs) -> Tuple[str, TemplateExecutionResults]:
 
             if line:
                 raise SQLTemplateException(f"{message.strip()} line {line[0]}")
-            else:
-                raise SQLTemplateException(f"{message.strip()}")
+            raise SQLTemplateException(f"{message.strip()}")
         except Exception as e:
             if isinstance(e, SQLTemplateException):
                 raise e
-            else:
-                logging.exception(f"Error on unbound local error: {e}")
-                raise ValueError(str(e))
+            logging.exception(f"Error on unbound local error: {e}")
+            raise ValueError(str(e))
     except TypeError as e:
         error = str(e)
         if "not supported between instances of 'Placeholder' and " in str(e):
@@ -1829,7 +1829,7 @@ def get_var_names(t: Template):
     >>> [v['name'] for v in get_var_names(Template("{% for name in secret_names %}{{tb_secret(name)}}{% end %}"))]
     ['secret_names', 'name', 'tb_secret', 'name']
 
-    Type casting with date function in elif chain (type is in reserved_vars):
+    Type casting with date function in elif chain:
     >>> t = Template("{% if kind == 'int' %}{{Int32(val, 0)}}{% elif kind == 'date' %}{{DateTime(val)}}{% elif kind == 'diff' %}{{date_diff_in_days(val, now)}}{% end %}")
     >>> [v['name'] for v in get_var_names(t)]
     ['kind', 'Int32', 'val', 'kind', 'DateTime', 'val', 'kind', 'date_diff_in_days', 'val', 'now']
@@ -1960,29 +1960,27 @@ def get_var_data(content, node_id=None):
         # ast.Constant nodes, so .value holds str/bytes/int/float/bool/None.
         if type(x) == ast.Constant:  # noqa: E721
             return x.value
-        elif type(x) == ast.Name:  # noqa: E721
+        if type(x) == ast.Name:  # noqa: E721
             return x.id
-        elif type(x) == ast.List:  # noqa: E721
+        if type(x) == ast.List:  # noqa: E721
             # List can hold different types
             return _get_list_var_data(x)
-        elif type(x) == ast.BinOp:  # noqa: E721
+        if type(x) == ast.BinOp:  # noqa: E721
             # in this case there could be several variables
             # if that's the case the left one is the main
             r = node_to_value(x.left)
             if not r:
                 r = node_to_value(x.right)
             return r
-        elif type(x) == ast.UnaryOp and type(x.operand) == ast.Constant:  # noqa: E721
+        if type(x) == ast.UnaryOp and type(x.operand) == ast.Constant:  # noqa: E721
             if type(x.op) == ast.USub:  # noqa: E721
                 return x.operand.value * -1
-            else:
-                return x.operand.value
-        else:
-            try:
-                return x.id
-            except Exception:
-                # don't let this ruin the parsing
-                pass
+            return x.operand.value
+        try:
+            return x.id
+        except Exception:
+            # don't let this ruin the parsing
+            pass
         return None
 
     def _get_list_var_data(x):
@@ -1994,7 +1992,7 @@ def get_var_data(content, node_id=None):
         # ast.Bytes/ast.Str/ast.Num/ast.NameConstant (removed in 3.14).
         if type(first_elem) == ast.Constant:  # noqa: E721
             return [elem.value for elem in x.elts]
-        elif type(first_elem) == ast.Name:  # noqa: E721
+        if type(first_elem) == ast.Name:  # noqa: E721
             return [elem.id for elem in x.elts]
 
         return []
@@ -2013,8 +2011,7 @@ def get_var_data(content, node_id=None):
     # to just `ast.parse(content)` once we confirm all environments use patched Python versions.
     def parse_content(content, retries=0):
         try:
-            parsed = ast.parse(content)
-            return parsed
+            return ast.parse(content)
         except Exception as e:
             if "AST constructor recursion depth mismatch" not in str(e):
                 raise e
@@ -2042,8 +2039,7 @@ def get_var_data(content, node_id=None):
                     if type(x) == ast.Call:  # noqa: E721
                         # Nested calls are traversed via `ast.iter_child_nodes` below.
                         continue
-                    else:
-                        args.append(node_to_value(x))
+                    args.append(node_to_value(x))
 
                 kwargs = {}
                 for x in node.keywords:
@@ -2385,7 +2381,24 @@ def wrap_vars(t, escape_arrays: bool = False):
     return var
 
 
-def get_used_tables_in_template(sql):
+def _constant_table_name(node: ast.expr) -> Optional[str]:
+    """Value of a constant table() argument: a literal or a concatenation of string literals."""
+    if isinstance(node, ast.Constant) and isinstance(node.value, (int, float)):
+        return str(node.value)
+    return _string_literal(node)
+
+
+def _string_literal(node: ast.expr) -> Optional[str]:
+    if isinstance(node, ast.Constant) and isinstance(node.value, str):
+        return node.value
+    if isinstance(node, ast.BinOp) and isinstance(node.op, ast.Add):
+        left, right = _string_literal(node.left), _string_literal(node.right)
+        if left is not None and right is not None:
+            return left + right
+    return None
+
+
+def get_used_tables_in_template(sql: str) -> List[str]:
     """
     >>> get_used_tables_in_template("select * from {{table('test')}}")
     ['test']
@@ -2399,30 +2412,42 @@ def get_used_tables_in_template(sql):
     []
     >>> get_used_tables_in_template("select * from {{table('my.test')}}, {{table('another.one')}}")
     ['my.test', 'another.one']
+
+    Only constant table names can be known without rendering, so anything else is skipped
+    >>> get_used_tables_in_template("select * from {{table(tbl)}}, {{TABLE('b')}}, {{table('my.' + 'c')}}")
+    ['b', 'my.c']
     """
+    return list(_get_used_tables_in_template_cached(sql))
+
+
+# Building the Template dominates the cost and the result only depends on the SQL, which deployments and
+# /v0/sql pipelines scan repeatedly
+@lru_cache(maxsize=2**10)
+def _get_used_tables_in_template_cached(sql: str) -> Tuple[str, ...]:
     try:
         t = Template(sql)
 
         def _n(chunks, tables):
             for x in chunks:
                 if type(x).__name__ == "_Expression":
-                    c = compile(x.expression, "<string>", "exec", dont_inherit=True)
-                    v = [x.lower() for x in c.co_names if x not in _namespace and x not in reserved_vars]
-                    if "table" in v:
-
-                        def _t(*args, **kwargs):
-                            return str(args[0])
-
-                        n = {"table": _t, "TABLE": _t}
-                        e = "_tt_tmp = %s" % x.expression
-                        exec_in(e, n)
-                        tables += [n["_tt_tmp"]]
+                    # Read the names from the AST instead of executing the expression: this runs outside
+                    # generate(), without the restricted builtins, on every pipe, query and deployment.
+                    for node in ast.walk(ast.parse(x.expression)):
+                        if (
+                            isinstance(node, ast.Call)
+                            and isinstance(node.func, ast.Name)
+                            and node.func.id in ("table", "TABLE")
+                            and node.args
+                        ):
+                            name = _constant_table_name(node.args[0])
+                            if name is not None:
+                                tables.append(name)
                 elif type(x).__name__ == "_ControlBlock":
                     _n(x.body.chunks, tables)
 
-        tables = []
+        tables: List[str] = []
         _n(t.file.body.chunks, tables)
-        return tables
+        return tuple(tables)
     except SecurityException as e:
         raise SQLTemplateException(e)
 
@@ -2531,8 +2556,7 @@ def format_SQLTemplateException_message(e: SQLTemplateException, vars_and_types:
             f"{REQUIRED_PARAM_NOT_DEFINED}. Check the parameters {join_with_different_last_separator(vars_with_default_none)}. Please provide a value or set a default value in the pipe code.",
             e.documentation,
         )
-    else:
-        raise e
+    raise e
 
 
 def render_sql_template(
@@ -2893,6 +2917,13 @@ def render_sql_template(
                 documentation="/cli/advanced-templates.html",
             )
         raise SQLTemplateException(str(e), documentation="/cli/advanced-templates.html")
+    except KeyError as e:
+        key = e.args[0] if e.args else "unknown"
+        raise SQLTemplateException(
+            f"Missing key '{key}' while rendering the template. "
+            "Ensure every object accessed in the template has this key.",
+            documentation="/cli/advanced-templates.html",
+        )
     except IndexError as e:
         # This happens when trying to access string indices on empty strings
         if "string index out of range" in str(e):
@@ -2906,12 +2937,12 @@ def render_sql_template(
         # https://gitlab.com/tinybird/analytics/-/issues/943
         if "length" in v and not v["length"]:
             raise SQLTemplateException("length cannot be used as a variable name or as a function inside of a template")
-        elif "missing 1 required positional argument" in str(e):
+        if "missing 1 required positional argument" in str(e):
             raise SQLTemplateException(
                 "one of the transform type functions is missing an argument",
                 documentation="/cli/advanced-templates.html#transform-types-functions",
             )
-        elif "not callable" in str(e) or "unhashable type" in str(e):
+        if "not callable" in str(e) or "unhashable type" in str(e):
             raise SQLTemplateException(
                 "wrong syntax, you might be using a not valid function inside a control block",
                 documentation="/cli/advanced-templates.html",
@@ -3056,7 +3087,7 @@ def render_template_with_secrets(
                     )
                 return '""'
             return value
-        elif default is not None:
+        if default is not None:
             if isinstance(default, str) and len(default) == 0:
                 if empty_secret_raises:
                     raise SQLTemplateException(
@@ -3064,10 +3095,9 @@ def render_template_with_secrets(
                     )
                 return '""'
             return default
-        else:
-            raise SQLTemplateException(
-                f"Cannot access secret '{secret_name}'. Check the secret exists in the Workspace and the token has the required scope."
-            )
+        raise SQLTemplateException(
+            f"Cannot access secret '{secret_name}'. Check the secret exists in the Workspace and the token has the required scope."
+        )
 
     # Create the template
     t = Template(content, name=name, autoescape=None)

tinybird/sql_template_fmt.py

@@ -33,8 +33,7 @@ class TBLine(Line):
         Tinybird => This is overriden from the base Line because we want SQL inside a template to be indented
         https://github.com/tconbeer/sqlfmt/blob/c11775b92d8a45f0e91d871b81a88a894d620bec/src/sqlfmt/line.py#L92
         """
-        prefix = INDENT * (self.depth[0] + self.depth[1])
-        return prefix
+        return INDENT * (self.depth[0] + self.depth[1])
 
 
 def from_nodes(
@@ -106,8 +105,7 @@ def _format_jinja_node(self, node: Node, max_length: int) -> bool:
 
         return tag.is_blackened
 
-    else:
-        return False
+    return False
 
 
 # Some monkey patching
@@ -158,8 +156,7 @@ class TinybirdNodeManager(NodeManager):
             TokenType.SET_OPERATOR,
         ):
             return " ".join(token.token.split())
-        else:
-            return token.token
+        return token.token
 
 
 class TinybirdDialect(ClickHouse):
@@ -217,8 +214,8 @@ class TinybirdDialect(ClickHouse):
             ],
         }
 
-        for section in override_rules:
-            for rule in override_rules[section]:
+        for section, rules in override_rules.items():
+            for rule in rules:
                 for rr in self.RULES[section]:
                     if rr.name == rule.name:
                         self.RULES[section].remove(rr)
@@ -230,12 +227,11 @@ class TinybirdDialect(ClickHouse):
         Creates and returns an analyzer that uses the Dialect's rules for lexing
         Custom NodeManager for Tinybird
         """
-        analyzer = Analyzer(
+        return Analyzer(
             line_length=line_length,
             rules=self.get_rules(),
             node_manager=TinybirdNodeManager(self.case_sensitive_names, self.lower_keywords),
         )
-        return analyzer
 
 
 @dataclass
@@ -253,11 +249,10 @@ class TBMode(Mode):
 def _calc_str(self) -> str:
     if self.is_multiline:
         return self.token.token + "\n"
-    else:
-        marker, comment_text = self._comment_parts()
-        if comment_text == "":
-            return marker + "\n"
-        return marker + " " + comment_text + "\n"
+    marker, comment_text = self._comment_parts()
+    if comment_text == "":
+        return marker + "\n"
+    return marker + " " + comment_text + "\n"
 
 
 Comment._calc_str = property(_calc_str)