diff options
| author | YDBot <[email protected]> | 2026-07-17 13:57:26 +0000 |
|---|---|---|
| committer | YDBot <[email protected]> | 2026-07-17 13:57:26 +0000 |
| commit | 462b35b1d3bc7a668b784cf0f313bec2793738e2 (patch) | |
| tree | 4e65b456392a49930301625b9320e74e01eeb0ee | |
| parent | 08a0f6a3ebbc624aa892835e03eeaa5c027f2bca (diff) | |
| parent | 880f1f4a67831e1f0732ea8d8d5a9ba43756d78e (diff) | |
Merge pull request #46914 from ydb-platform/merge-rightlib-260717-1016
148 files changed, 3977 insertions, 1166 deletions
diff --git a/build/mapping.conf.json b/build/mapping.conf.json index 6483820b6db..cef8ab62008 100644 --- a/build/mapping.conf.json +++ b/build/mapping.conf.json @@ -622,6 +622,7 @@ "12727426752": "{registry_endpoint}/12727426752", "12765030473": "{registry_endpoint}/12765030473", "12778660629": "{registry_endpoint}/12778660629", + "12807488115": "{registry_endpoint}/12807488115", "5486713852": "{registry_endpoint}/5486713852", "5514352253": "{registry_endpoint}/5514352253", "5523579199": "{registry_endpoint}/5523579199", @@ -2109,6 +2110,7 @@ "12606129409": "{registry_endpoint}/12606129409", "12636194389": "{registry_endpoint}/12636194389", "12666249275": "{registry_endpoint}/12666249275", + "12811590910": "{registry_endpoint}/12811590910", "6662972834": "{registry_endpoint}/6662972834", "11823170312": "{registry_endpoint}/11823170312", "11879814525": "{registry_endpoint}/11879814525", @@ -2129,6 +2131,7 @@ "12606127844": "{registry_endpoint}/12606127844", "12636192320": "{registry_endpoint}/12636192320", "12666248010": "{registry_endpoint}/12666248010", + "12811589861": "{registry_endpoint}/12811589861", "6662965433": "{registry_endpoint}/6662965433", "11823173798": "{registry_endpoint}/11823173798", "11879819941": "{registry_endpoint}/11879819941", @@ -2149,6 +2152,7 @@ "12606130988": "{registry_endpoint}/12606130988", "12636196166": "{registry_endpoint}/12636196166", "12666251081": "{registry_endpoint}/12666251081", + "12811592057": "{registry_endpoint}/12811592057", "6662980150": "{registry_endpoint}/6662980150", "5476908047": "{registry_endpoint}/5476908047", "5509380757": "{registry_endpoint}/5509380757", @@ -3040,6 +3044,7 @@ "12727426752": "devtools/ya/test/programs/test_tool/bin/test_tool for linux", "12765030473": "devtools/ya/test/programs/test_tool/bin/test_tool for linux", "12778660629": "devtools/ya/test/programs/test_tool/bin/test_tool for linux", + "12807488115": "devtools/ya/test/programs/test_tool/bin/test_tool for linux", "5486713852": "devtools/ya/test/programs/test_tool/bin/test_tool for linux", "5514352253": "devtools/ya/test/programs/test_tool/bin/test_tool for linux", "5523579199": "devtools/ya/test/programs/test_tool/bin/test_tool for linux", @@ -4527,6 +4532,7 @@ "12606129409": "yfm-docs for darwin", "12636194389": "yfm-docs for darwin", "12666249275": "yfm-docs for darwin", + "12811590910": "yfm-docs for darwin", "6662972834": "yfm-docs for darwin", "11823170312": "yfm-docs for linux", "11879814525": "yfm-docs for linux", @@ -4547,6 +4553,7 @@ "12606127844": "yfm-docs for linux", "12636192320": "yfm-docs for linux", "12666248010": "yfm-docs for linux", + "12811589861": "yfm-docs for linux", "6662965433": "yfm-docs for linux", "11823173798": "yfm-docs for win32", "11879819941": "yfm-docs for win32", @@ -4567,6 +4574,7 @@ "12606130988": "yfm-docs for win32", "12636196166": "yfm-docs for win32", "12666251081": "yfm-docs for win32", + "12811592057": "yfm-docs for win32", "6662980150": "yfm-docs for win32", "5476908047": "ymake for darwin", "5509380757": "ymake for darwin", diff --git a/build/platform/test_tool/host.ya.make.inc b/build/platform/test_tool/host.ya.make.inc index e23ae9ddf1a..45904018c0f 100644 --- a/build/platform/test_tool/host.ya.make.inc +++ b/build/platform/test_tool/host.ya.make.inc @@ -1,12 +1,12 @@ IF (HOST_OS_DARWIN AND HOST_ARCH_X86_64) - DECLARE_EXTERNAL_RESOURCE(TEST_TOOL_HOST sbr:12778677593) + DECLARE_EXTERNAL_RESOURCE(TEST_TOOL_HOST sbr:12807476180) ELSEIF (HOST_OS_DARWIN AND HOST_ARCH_ARM64) - DECLARE_EXTERNAL_RESOURCE(TEST_TOOL_HOST sbr:12778675385) + DECLARE_EXTERNAL_RESOURCE(TEST_TOOL_HOST sbr:12807473502) ELSEIF (HOST_OS_LINUX AND HOST_ARCH_X86_64) - DECLARE_EXTERNAL_RESOURCE(TEST_TOOL_HOST sbr:12778682489) + DECLARE_EXTERNAL_RESOURCE(TEST_TOOL_HOST sbr:12807481043) ELSEIF (HOST_OS_LINUX AND HOST_ARCH_AARCH64) - DECLARE_EXTERNAL_RESOURCE(TEST_TOOL_HOST sbr:12778673354) + DECLARE_EXTERNAL_RESOURCE(TEST_TOOL_HOST sbr:12807470623) ELSEIF (HOST_OS_WINDOWS AND HOST_ARCH_X86_64) - DECLARE_EXTERNAL_RESOURCE(TEST_TOOL_HOST sbr:12778679810) + DECLARE_EXTERNAL_RESOURCE(TEST_TOOL_HOST sbr:12807478544) ENDIF() diff --git a/build/platform/test_tool/host_os.ya.make.inc b/build/platform/test_tool/host_os.ya.make.inc index d65679840ec..f6061502b1c 100644 --- a/build/platform/test_tool/host_os.ya.make.inc +++ b/build/platform/test_tool/host_os.ya.make.inc @@ -1,12 +1,12 @@ IF (HOST_OS_DARWIN AND HOST_ARCH_X86_64) - DECLARE_EXTERNAL_RESOURCE(TEST_TOOL_HOST sbr:12778655677) + DECLARE_EXTERNAL_RESOURCE(TEST_TOOL_HOST sbr:12807484254) ELSEIF (HOST_OS_DARWIN AND HOST_ARCH_ARM64) - DECLARE_EXTERNAL_RESOURCE(TEST_TOOL_HOST sbr:12778653112) + DECLARE_EXTERNAL_RESOURCE(TEST_TOOL_HOST sbr:12807482262) ELSEIF (HOST_OS_LINUX AND HOST_ARCH_X86_64) - DECLARE_EXTERNAL_RESOURCE(TEST_TOOL_HOST sbr:12778660629) + DECLARE_EXTERNAL_RESOURCE(TEST_TOOL_HOST sbr:12807488115) ELSEIF (HOST_OS_LINUX AND HOST_ARCH_AARCH64) - DECLARE_EXTERNAL_RESOURCE(TEST_TOOL_HOST sbr:12778650528) + DECLARE_EXTERNAL_RESOURCE(TEST_TOOL_HOST sbr:12807480382) ELSEIF (HOST_OS_WINDOWS AND HOST_ARCH_X86_64) - DECLARE_EXTERNAL_RESOURCE(TEST_TOOL_HOST sbr:12778657924) + DECLARE_EXTERNAL_RESOURCE(TEST_TOOL_HOST sbr:12807486192) ENDIF() diff --git a/build/platform/yfm/resources.json b/build/platform/yfm/resources.json index 08b886b25aa..0b90b6fab0a 100644 --- a/build/platform/yfm/resources.json +++ b/build/platform/yfm/resources.json @@ -1,16 +1,16 @@ { "by_platform": { "win32-x86_64": { - "uri": "sbr:12666251081" + "uri": "sbr:12811592057" }, "darwin-x86_64": { - "uri": "sbr:12666249275" + "uri": "sbr:12811590910" }, "linux-x86_64": { - "uri": "sbr:12666248010" + "uri": "sbr:12811589861" }, "darwin-arm64": { - "uri": "sbr:12666249275" + "uri": "sbr:12811590910" } } } diff --git a/contrib/python/argcomplete/py3/.dist-info/METADATA b/contrib/python/argcomplete/py3/.dist-info/METADATA index 33d51a548ba..fa5e1d45d3a 100644 --- a/contrib/python/argcomplete/py3/.dist-info/METADATA +++ b/contrib/python/argcomplete/py3/.dist-info/METADATA @@ -1,6 +1,6 @@ Metadata-Version: 2.4 Name: argcomplete -Version: 3.6.3 +Version: 3.7.0 Summary: Bash tab completion for argparse Project-URL: Homepage, https://github.com/kislyuk/argcomplete Project-URL: Documentation, https://kislyuk.github.io/argcomplete @@ -22,12 +22,12 @@ Classifier: Operating System :: MacOS :: MacOS X Classifier: Operating System :: POSIX Classifier: Programming Language :: Python Classifier: Programming Language :: Python :: 3 -Classifier: Programming Language :: Python :: 3.8 Classifier: Programming Language :: Python :: 3.9 Classifier: Programming Language :: Python :: 3.10 Classifier: Programming Language :: Python :: 3.11 Classifier: Programming Language :: Python :: 3.12 Classifier: Programming Language :: Python :: 3.13 +Classifier: Programming Language :: Python :: 3.14 Classifier: Programming Language :: Python :: Implementation :: CPython Classifier: Programming Language :: Python :: Implementation :: PyPy Classifier: Topic :: Software Development diff --git a/contrib/python/argcomplete/py3/.dist-info/entry_points.txt b/contrib/python/argcomplete/py3/.dist-info/entry_points.txt index f13e435fafb..874a9e1481d 100644 --- a/contrib/python/argcomplete/py3/.dist-info/entry_points.txt +++ b/contrib/python/argcomplete/py3/.dist-info/entry_points.txt @@ -1,4 +1,3 @@ [console_scripts] activate-global-python-argcomplete = argcomplete.scripts.activate_global_python_argcomplete:main -python-argcomplete-check-easy-install-script = argcomplete.scripts.python_argcomplete_check_easy_install_script:main register-python-argcomplete = argcomplete.scripts.register_python_argcomplete:main diff --git a/contrib/python/argcomplete/py3/argcomplete/bash_completion.d/_python-argcomplete b/contrib/python/argcomplete/py3/argcomplete/bash_completion.d/_python-argcomplete index 1155f1ce5e8..6d395c66ed9 100644 --- a/contrib/python/argcomplete/py3/argcomplete/bash_completion.d/_python-argcomplete +++ b/contrib/python/argcomplete/py3/argcomplete/bash_completion.d/_python-argcomplete @@ -190,10 +190,7 @@ _python_argcomplete_global() { interpreter=($interpreter) fi - if (__python_argcomplete_scan_head_noerr "$SCRIPT_NAME" easy_install \ - && "${interpreter[@]}" "$(__python_argcomplete_which python-argcomplete-check-easy-install-script)" "$SCRIPT_NAME") >/dev/null 2>&1; then - ARGCOMPLETE=1 - elif ([[ "${interpreter[@]}" == *python* ]] || [[ "${interpreter[@]}" == *pypy* ]])\ + if ([[ "${interpreter[@]}" == *python* ]] || [[ "${interpreter[@]}" == *pypy* ]])\ && __python_argcomplete_run "${interpreter[@]}" -m argcomplete._check_console_script "$SCRIPT_NAME"; then ARGCOMPLETE=1 fi diff --git a/contrib/python/argcomplete/py3/argcomplete/completers.py b/contrib/python/argcomplete/py3/argcomplete/completers.py index 4c01e695189..48715811898 100644 --- a/contrib/python/argcomplete/py3/argcomplete/completers.py +++ b/contrib/python/argcomplete/py3/argcomplete/completers.py @@ -4,6 +4,7 @@ import argparse import os import subprocess +from shlex import quote def _call(*args, **kwargs): @@ -64,20 +65,22 @@ class FilesCompleter(BaseCompleter): # correctly in older versions and calling bind makes them available. For details, see # https://savannah.gnu.org/support/index.php?111125 files = _call( - ["bash", "-c", "bind; compgen -A directory -- '{p}'".format(p=prefix)], stderr=subprocess.DEVNULL + ["bash", "-c", "bind; compgen -A directory -- {p}".format(p=quote(prefix))], + stderr=subprocess.DEVNULL, ) completion += [f + "/" for f in files] for x in self.allowednames: completion += _call( - ["bash", "-c", "bind; compgen -A file -X '!*.{0}' -- '{p}'".format(x, p=prefix)], + ["bash", "-c", "bind; compgen -A file -X '!*.{0}' -- {p}".format(x, p=quote(prefix))], stderr=subprocess.DEVNULL, ) else: completion += _call( - ["bash", "-c", "bind; compgen -A file -- '{p}'".format(p=prefix)], stderr=subprocess.DEVNULL + ["bash", "-c", "bind; compgen -A file -- {p}".format(p=quote(prefix))], stderr=subprocess.DEVNULL ) anticomp = _call( - ["bash", "-c", "bind; compgen -A directory -- '{p}'".format(p=prefix)], stderr=subprocess.DEVNULL + ["bash", "-c", "bind; compgen -A directory -- {p}".format(p=quote(prefix))], + stderr=subprocess.DEVNULL, ) completion = list(set(completion) - set(anticomp)) diff --git a/contrib/python/argcomplete/py3/argcomplete/finders.py b/contrib/python/argcomplete/py3/argcomplete/finders.py index ab8bb1d90a2..2159b7dea2f 100644 --- a/contrib/python/argcomplete/py3/argcomplete/finders.py +++ b/contrib/python/argcomplete/py3/argcomplete/finders.py @@ -527,7 +527,7 @@ class CompletionFinder(object): # (extended to characters other than the colon). if last_wordbreak_pos is not None: completions = [c[last_wordbreak_pos + 1 :] for c in completions] - special_chars += "();<>|&!`$* \t\n\"'" + special_chars += "();<>|&!`$*?[]{} \t\n\"'" elif cword_prequote == '"': special_chars += '"`$!' diff --git a/contrib/python/argcomplete/py3/argcomplete/packages/_argparse.py b/contrib/python/argcomplete/py3/argcomplete/packages/_argparse.py index f6ecb1f4172..574d9d8207a 100644 --- a/contrib/python/argcomplete/py3/argcomplete/packages/_argparse.py +++ b/contrib/python/argcomplete/py3/argcomplete/packages/_argparse.py @@ -19,7 +19,13 @@ from argparse import ( _SubParsersAction, ) from gettext import gettext -from typing import Dict, List, Set, Tuple +from typing import Dict, List, Optional, Set, Tuple, Union, cast + +_OptionTuple = Union[ + Tuple[Optional[Action], str, Optional[str]], + Tuple[Optional[Action], str, Optional[str], Optional[str]], +] +_OptionTupleEntry = Union[_OptionTuple, List[_OptionTuple]] _num_consumed_args: Dict[Action, int] = {} @@ -97,7 +103,7 @@ class IntrospectiveArgumentParser(ArgumentParser): # find all option indices, and determine the arg_string_pattern # which has an 'O' if there is an option at an index, # an 'A' if there is an argument, or a '-' if there is a '--' - option_string_indices = {} + option_string_indices: Dict[int, _OptionTupleEntry] = {} arg_string_pattern_parts = [] arg_strings_iter = iter(arg_strings) for i, arg_string in enumerate(arg_strings_iter): @@ -114,7 +120,7 @@ class IntrospectiveArgumentParser(ArgumentParser): if option_tuple is None: pattern = 'A' else: - option_string_indices[i] = option_tuple + option_string_indices[i] = cast(_OptionTupleEntry, option_tuple) pattern = 'O' arg_string_pattern_parts.append(pattern) @@ -161,9 +167,11 @@ class IntrospectiveArgumentParser(ArgumentParser): # function to convert arg_strings into an optional action def consume_optional(start_index): # get the optional identified at this index - option_tuple = option_string_indices[start_index] - if isinstance(option_tuple, list): # Python 3.12.7+ - option_tuple = option_tuple[0] + raw_option_tuple = option_string_indices[start_index] + if isinstance(raw_option_tuple, list): # Python 3.12.7+ + option_tuple = raw_option_tuple[0] + else: + option_tuple = raw_option_tuple if len(option_tuple) == 3: action, option_string, explicit_arg = option_tuple else: # Python 3.11.9+, 3.12.3+, 3.13+ @@ -241,8 +249,8 @@ class IntrospectiveArgumentParser(ArgumentParser): # add the Optional to the list and return the index at which # the Optional's string args stopped assert action_tuples - for action, args, option_string in action_tuples: - take_action(action, args, option_string) + for optional_action, args, option_string in action_tuples: + take_action(optional_action, args, option_string) return stop # the list of Positionals left to be parsed; this is modified diff --git a/contrib/python/argcomplete/py3/argcomplete/scripts/python_argcomplete_check_easy_install_script.py b/contrib/python/argcomplete/py3/argcomplete/scripts/python_argcomplete_check_easy_install_script.py deleted file mode 100644 index 0eb744c9e96..00000000000 --- a/contrib/python/argcomplete/py3/argcomplete/scripts/python_argcomplete_check_easy_install_script.py +++ /dev/null @@ -1,84 +0,0 @@ -#!/usr/bin/env python3 - -# Copyright 2012-2023, Andrey Kislyuk and argcomplete contributors. -# Licensed under the Apache License. See https://github.com/kislyuk/argcomplete for more info. - -""" -This script is part of the Python argcomplete package (https://github.com/kislyuk/argcomplete). -It is used to check if an EASY-INSTALL-SCRIPT wrapper redirects to a script that contains the string -"PYTHON_ARGCOMPLETE_OK". If you have enabled global completion in argcomplete, the completion hook will run it every -time you press <TAB> in your shell. - -Usage: - python-argcomplete-check-easy-install-script <input executable file> -""" - -import sys - -# PEP 366 -__package__ = "argcomplete.scripts" - - -def main(): - if len(sys.argv) != 2: - sys.exit(__doc__) - - sys.tracebacklimit = 0 - - with open(sys.argv[1]) as fh: - line1, head = fh.read(1024).split("\n", 1)[:2] - if line1.startswith("#") and ("py" in line1 or "Py" in line1): - import re - - lines = head.split("\n", 12) - for line in lines: - if line.startswith("# EASY-INSTALL-SCRIPT"): - import pkg_resources # type: ignore - - re_match = re.match("# EASY-INSTALL-SCRIPT: '(.+)','(.+)'", line) - assert re_match is not None - dist, script = re_match.groups() - if "PYTHON_ARGCOMPLETE_OK" in pkg_resources.get_distribution(dist).get_metadata( - "scripts/" + script - ): - return 0 - elif line.startswith("# EASY-INSTALL-ENTRY-SCRIPT"): - re_match = re.match("# EASY-INSTALL-ENTRY-SCRIPT: '(.+)','(.+)','(.+)'", line) - assert re_match is not None - dist, group, name = re_match.groups() - import pkgutil - - import pkg_resources # type: ignore - - entry_point_info = pkg_resources.get_distribution(dist).get_entry_info(group, name) - assert entry_point_info is not None - module_name = entry_point_info.module_name - with open(pkgutil.get_loader(module_name).get_filename()) as mod_fh: # type: ignore - if "PYTHON_ARGCOMPLETE_OK" in mod_fh.read(1024): - return 0 - elif line.startswith("# EASY-INSTALL-DEV-SCRIPT"): - for line2 in lines: - if line2.startswith("__file__"): - re_match = re.match("__file__ = '(.+)'", line2) - assert re_match is not None - filename = re_match.group(1) - with open(filename) as mod_fh: - if "PYTHON_ARGCOMPLETE_OK" in mod_fh.read(1024): - return 0 - elif line.startswith("# PBR Generated"): - re_match = re.search("from (.*) import", head) - assert re_match is not None - module = re_match.groups()[0] - import pkgutil - - import pkg_resources # type: ignore - - with open(pkgutil.get_loader(module).get_filename()) as mod_fh: # type: ignore - if "PYTHON_ARGCOMPLETE_OK" in mod_fh.read(1024): - return 0 - - return 1 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/contrib/python/argcomplete/py3/ya.make b/contrib/python/argcomplete/py3/ya.make index 1ffd0082bb2..397e1ff3f2e 100644 --- a/contrib/python/argcomplete/py3/ya.make +++ b/contrib/python/argcomplete/py3/ya.make @@ -2,7 +2,7 @@ PY3_LIBRARY() -VERSION(3.6.3) +VERSION(3.7.0) LICENSE(Apache-2.0) @@ -23,7 +23,6 @@ PY_SRCS( argcomplete/packages/_shlex.py argcomplete/scripts/__init__.py argcomplete/scripts/activate_global_python_argcomplete.py - argcomplete/scripts/python_argcomplete_check_easy_install_script.py argcomplete/scripts/register_python_argcomplete.py argcomplete/shell_integration.py ) diff --git a/contrib/python/clickhouse-connect/.dist-info/METADATA b/contrib/python/clickhouse-connect/.dist-info/METADATA index cf511761a8d..656e95416df 100644 --- a/contrib/python/clickhouse-connect/.dist-info/METADATA +++ b/contrib/python/clickhouse-connect/.dist-info/METADATA @@ -1,6 +1,6 @@ Metadata-Version: 2.4 Name: clickhouse-connect -Version: 1.3.0 +Version: 1.4.1 Summary: ClickHouse Database Core Driver for Python, Pandas, and Superset Home-page: https://github.com/ClickHouse/clickhouse-connect Author: ClickHouse Inc. diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/_version.py b/contrib/python/clickhouse-connect/clickhouse_connect/_version.py index d28b3ddc3b1..6abaa204304 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/_version.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/_version.py @@ -1 +1 @@ -version = "1.3.0" +version = "1.4.1" diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/alembic/adapter.py b/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/alembic/adapter.py index d444b69e73a..0319a329855 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/alembic/adapter.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/alembic/adapter.py @@ -5,6 +5,7 @@ from alembic.operations import Operations, ops from alembic.runtime.migration import MigrationContext from alembic.util import DispatchPriority, PriorityDispatchResult +from clickhouse_connect.cc_sqlalchemy.alembic.impl import ClickHouseImpl from clickhouse_connect.cc_sqlalchemy.datatypes.base import ChSqlaType from clickhouse_connect.cc_sqlalchemy.sql.ddlcompiler import ClickHouseDDLHelper @@ -57,6 +58,12 @@ def clickhouse_writer(context: MigrationContext, revision, directives): _add_common_imports(directive) +def _is_clickhouse_autogen(autogen_context: AutogenContext) -> bool: + """True only when the active migration context targets the ClickHouse dialect.""" + migration_context = getattr(autogen_context, "migration_context", None) + return isinstance(getattr(migration_context, "impl", None), ClickHouseImpl) + + def render_clickhouse_column(column, autogen_context: AutogenContext) -> str: rendered = render._user_defined_render("column", column, autogen_context) if rendered is not False: @@ -98,8 +105,19 @@ def render_clickhouse_column(column, autogen_context: AutogenContext) -> str: ) +# Alembic renderers have no dialect qualifier, so replace=True overrides rendering +# process-wide. Capture each built-in renderer before replacing it and delegate to it for +# non-ClickHouse dialects so autogenerate stays correct for other databases (#832). Held on +# the module so importlib.reload does not re-capture one of our own renderers and recurse. +_DEFAULT_RENDERERS = globals().get("_DEFAULT_RENDERERS") or { + op: render.renderers.dispatch(op) for op in (ops.CreateTableOp, ops.AddColumnOp, ops.DropTableOp) +} + + @render.renderers.dispatch_for(ops.CreateTableOp, replace=True) def render_create_table(autogen_context: AutogenContext, op: ops.CreateTableOp) -> str: + if not _is_clickhouse_autogen(autogen_context): + return _DEFAULT_RENDERERS[ops.CreateTableOp](autogen_context, op) table = op.to_table() args = [column for column in [render_clickhouse_column(column, autogen_context) for column in table.columns] if column] + sorted( @@ -138,6 +156,8 @@ def render_create_table(autogen_context: AutogenContext, op: ops.CreateTableOp) @render.renderers.dispatch_for(ops.AddColumnOp, replace=True) def render_add_column(autogen_context: AutogenContext, op: ops.AddColumnOp) -> str: + if not _is_clickhouse_autogen(autogen_context): + return _DEFAULT_RENDERERS[ops.AddColumnOp](autogen_context, op) schema, table_name, column, if_not_exists = op.schema, op.table_name, op.column, op.if_not_exists prefix = render._alembic_autogenerate_prefix(autogen_context) rendered_column = render_clickhouse_column(column, autogen_context) @@ -155,6 +175,8 @@ def render_add_column(autogen_context: AutogenContext, op: ops.AddColumnOp) -> s @render.renderers.dispatch_for(ops.DropTableOp, replace=True) def render_drop_table(autogen_context: AutogenContext, op: ops.DropTableOp) -> str: + if not _is_clickhouse_autogen(autogen_context): + return _DEFAULT_RENDERERS[ops.DropTableOp](autogen_context, op) prefix = render._alembic_autogenerate_prefix(autogen_context) rendered = f"{prefix}drop_table({render._ident(op.table_name)!r}" arguments = [] diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/alembic/impl.py b/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/alembic/impl.py index 1f942979c07..32e2b985cd4 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/alembic/impl.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/alembic/impl.py @@ -65,7 +65,7 @@ class ClickHouseImpl(DefaultImpl): if self.connection is None: return try: - self.connection.connection.driver_connection.client._add_integration_tag("alembic") + self.connection.connection.driver_connection.client._add_integration_tag("alembic") # type: ignore[union-attr] except Exception: pass @@ -100,7 +100,7 @@ class ClickHouseImpl(DefaultImpl): construct, execution_options=execution_options, multiparams=multiparams, - params=params or {}, + params=params or {}, # type: ignore[arg-type] ) def add_column( diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/alembic/utils.py b/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/alembic/utils.py index 35ff8fb7b94..14ac01c27fc 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/alembic/utils.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/alembic/utils.py @@ -75,9 +75,9 @@ def prevent_empty_migrations(writer_fn: Callable) -> Callable: if not directives: return config = context.config - if getattr(config.cmd_opts, "autogenerate", False): + if config is not None and getattr(config.cmd_opts, "autogenerate", False): script = directives[0] - if script.upgrade_ops.is_empty(): + if script.upgrade_ops is not None and script.upgrade_ops.is_empty(): directives.clear() return writer_fn(context, revision, directives) diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/datatypes/base.py b/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/datatypes/base.py index a4c1d4a94f2..175c203668e 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/datatypes/base.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/datatypes/base.py @@ -15,11 +15,10 @@ class ChSqlaType: subclasses will inherit from TypeEngine. """ - ch_type: ClickHouseType = None + ch_type: ClickHouseType | None = None generic_type: None - _ch_type_cls = None - _instance = None - _instance_cache: dict[TypeDef, "ChSqlaType"] = None + _ch_type_cls: type[ClickHouseType] | None = None + _instance_cache: dict[TypeDef, "ChSqlaType"] | None = None def __init_subclass__(cls): """ @@ -44,7 +43,7 @@ class ChSqlaType: :param type_def: -- TypeDef tuple that defines arguments for this instance :return: Shared instance of a configured ChSqlaType """ - return cls._instance_cache.setdefault(type_def, cls(type_def=type_def)) + return cls._instance_cache.setdefault(type_def, cls(type_def=type_def)) # type: ignore[union-attr] def __init__(self, type_def: TypeDef = EMPTY_TYPE_DEF): """ @@ -55,7 +54,7 @@ class ChSqlaType: parse_name function """ self.type_def = type_def - self.ch_type = self._ch_type_cls.build(type_def) + self.ch_type = self._ch_type_cls.build(type_def) # type: ignore[union-attr] @property def name(self): @@ -124,7 +123,7 @@ class CaseInsensitiveDict(dict): sqla_type_map: dict[str, type[ChSqlaType]] = CaseInsensitiveDict() -schema_types = [] +schema_types: list[str] = [] def sqla_type_from_name(name: str) -> ChSqlaType: diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/datatypes/sqltypes.py b/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/datatypes/sqltypes.py index 5356150545b..ad7ee5ef173 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/datatypes/sqltypes.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/datatypes/sqltypes.py @@ -2,6 +2,7 @@ import ipaddress import uuid from collections.abc import Sequence from enum import Enum as PyEnum +from typing import Any, cast from sqlalchemy.exc import ArgumentError from sqlalchemy.types import ( @@ -10,6 +11,7 @@ from sqlalchemy.types import ( Integer, Interval, Numeric, + TypeEngine, UserDefinedType, ) from sqlalchemy.types import ( @@ -33,67 +35,67 @@ from clickhouse_connect.driver import tzutil from clickhouse_connect.driver.common import decimal_prec -class Int8(ChSqlaType, Integer): +class Int8(ChSqlaType, Integer): # type: ignore[misc] pass -class UInt8(ChSqlaType, Integer): +class UInt8(ChSqlaType, Integer): # type: ignore[misc] pass -class Int16(ChSqlaType, Integer): +class Int16(ChSqlaType, Integer): # type: ignore[misc] pass -class UInt16(ChSqlaType, Integer): +class UInt16(ChSqlaType, Integer): # type: ignore[misc] pass -class Int32(ChSqlaType, Integer): +class Int32(ChSqlaType, Integer): # type: ignore[misc] pass -class UInt32(ChSqlaType, Integer): +class UInt32(ChSqlaType, Integer): # type: ignore[misc] pass -class Int64(ChSqlaType, Integer): +class Int64(ChSqlaType, Integer): # type: ignore[misc] pass -class UInt64(ChSqlaType, Integer): +class UInt64(ChSqlaType, Integer): # type: ignore[misc] pass -class Int128(ChSqlaType, Integer): +class Int128(ChSqlaType, Integer): # type: ignore[misc] pass -class UInt128(ChSqlaType, Integer): +class UInt128(ChSqlaType, Integer): # type: ignore[misc] pass -class Int256(ChSqlaType, Integer): +class Int256(ChSqlaType, Integer): # type: ignore[misc] pass -class UInt256(ChSqlaType, Integer): +class UInt256(ChSqlaType, Integer): # type: ignore[misc] pass -class Float32(ChSqlaType, Float): +class Float32(ChSqlaType, Float): # type: ignore[misc] def __init__(self, type_def: TypeDef = EMPTY_TYPE_DEF): ChSqlaType.__init__(self, type_def) Float.__init__(self) -class Float64(ChSqlaType, Float): +class Float64(ChSqlaType, Float): # type: ignore[misc] def __init__(self, type_def: TypeDef = EMPTY_TYPE_DEF): ChSqlaType.__init__(self, type_def) Float.__init__(self) -class Bool(ChSqlaType, SqlaBoolean): +class Bool(ChSqlaType, SqlaBoolean): # type: ignore[misc] def __init__(self, type_def: TypeDef = EMPTY_TYPE_DEF, **kwargs): ChSqlaType.__init__(self, type_def) SqlaBoolean.__init__(self, **kwargs) @@ -103,10 +105,10 @@ class Boolean(Bool): pass -class Decimal(ChSqlaType, Numeric): +class Decimal(ChSqlaType, Numeric): # type: ignore[misc] dec_size = 0 - def __init__(self, precision: int = 0, scale: int = 0, type_def: TypeDef = None): + def __init__(self, precision: int = 0, scale: int = 0, type_def: TypeDef | None = None): """ Construct either with precision and scale (for DDL), or a TypeDef with those values (by name) :param precision: Number of digits the Decimal @@ -143,16 +145,16 @@ class Decimal256(Decimal): dec_size = 256 -class Enum(ChSqlaType, UserDefinedType): +class Enum(ChSqlaType, UserDefinedType): # type: ignore[misc] _size = 16 python_type = str def __init__( self, - enum: type[PyEnum] = None, - keys: Sequence[str] = None, - values: Sequence[int] = None, - type_def: TypeDef = None, + enum: type[PyEnum] | None = None, + keys: Sequence[str] | None = None, + values: Sequence[int] | None = None, + type_def: TypeDef | None = None, ): """ Construct a ClickHouse enum either from a Python Enum or parallel lists of keys and value. Note that @@ -166,6 +168,8 @@ class Enum(ChSqlaType, UserDefinedType): if enum: keys = [e.name for e in enum] values = [e.value for e in enum] + if keys is None or values is None: + raise ArgumentError("Enum requires either a Python enum or both 'keys' and 'values'") self._validate(keys, values) if self.__class__.__name__ == "Enum": if max(values) <= 127 and min(values) >= -128: @@ -176,7 +180,7 @@ class Enum(ChSqlaType, UserDefinedType): super().__init__(type_def) @classmethod - def _validate(cls, keys: Sequence, values: Sequence): + def _validate(cls, keys: Sequence[str], values: Sequence[int]): bad_key = next((x for x in keys if not isinstance(x, str)), None) if bad_key: raise ArgumentError(f"ClickHouse enum key {bad_key} is not a string") @@ -199,63 +203,63 @@ class Enum16(Enum): _ch_type_cls = ChEnum16 -class String(ChSqlaType, UserDefinedType): +class String(ChSqlaType, UserDefinedType): # type: ignore[misc] python_type = str -class FixedString(ChSqlaType, SqlaString): - def __init__(self, size: int = -1, type_def: TypeDef = None): +class FixedString(ChSqlaType, SqlaString): # type: ignore[misc] + def __init__(self, size: int = -1, type_def: TypeDef | None = None): if not type_def: type_def = TypeDef(values=(size,)) ChSqlaType.__init__(self, type_def) SqlaString.__init__(self, size) -class IPv4(ChSqlaType, UserDefinedType): +class IPv4(ChSqlaType, UserDefinedType): # type: ignore[misc] python_type = ipaddress.IPv4Address -class IPv6(ChSqlaType, UserDefinedType): +class IPv6(ChSqlaType, UserDefinedType): # type: ignore[misc] python_type = ipaddress.IPv6Address -class UUID(ChSqlaType, UserDefinedType): +class UUID(ChSqlaType, UserDefinedType): # type: ignore[misc] python_type = uuid.UUID -class Nothing(ChSqlaType, UserDefinedType): +class Nothing(ChSqlaType, UserDefinedType): # type: ignore[misc] python_type = type(None) -class Point(ChSqlaType, UserDefinedType): +class Point(ChSqlaType, UserDefinedType): # type: ignore[misc] python_type = tuple -class Ring(ChSqlaType, UserDefinedType): +class Ring(ChSqlaType, UserDefinedType): # type: ignore[misc] python_type = list -class Polygon(ChSqlaType, UserDefinedType): +class Polygon(ChSqlaType, UserDefinedType): # type: ignore[misc] python_type = list -class MultiPolygon(ChSqlaType, UserDefinedType): +class MultiPolygon(ChSqlaType, UserDefinedType): # type: ignore[misc] python_type = list -class LineString(ChSqlaType, UserDefinedType): +class LineString(ChSqlaType, UserDefinedType): # type: ignore[misc] python_type = list -class MultiLineString(ChSqlaType, UserDefinedType): +class MultiLineString(ChSqlaType, UserDefinedType): # type: ignore[misc] python_type = list -class Date(ChSqlaType, SqlaDate): +class Date(ChSqlaType, SqlaDate): # type: ignore[misc] pass -class Date32(ChSqlaType, SqlaDate): +class Date32(ChSqlaType, SqlaDate): # type: ignore[misc] pass @@ -283,8 +287,8 @@ def _resolve_tz_alias(tz, timezone): return tz -class DateTime(ChSqlaType, SqlaDateTime): - def __init__(self, tz: str = None, type_def: TypeDef = None, timezone=_TIMEZONE_SENTINEL): +class DateTime(ChSqlaType, SqlaDateTime): # type: ignore[misc] + def __init__(self, tz: str | None = None, type_def: TypeDef | None = None, timezone=_TIMEZONE_SENTINEL): """tz / timezone: IANA zone string (resolved via zoneinfo; install `tzdata` on Windows).""" tz = _resolve_tz_alias(tz, timezone) if not type_def: @@ -297,8 +301,8 @@ class DateTime(ChSqlaType, SqlaDateTime): SqlaDateTime.__init__(self) -class DateTime64(ChSqlaType, SqlaDateTime): - def __init__(self, precision: int = None, tz: str = None, type_def: TypeDef = None, timezone=_TIMEZONE_SENTINEL): +class DateTime64(ChSqlaType, SqlaDateTime): # type: ignore[misc] + def __init__(self, precision: int | None = None, tz: str | None = None, type_def: TypeDef | None = None, timezone=_TIMEZONE_SENTINEL): """precision: 3/6/9 for ms/us/ns. tz / timezone: IANA zone string.""" tz = _resolve_tz_alias(tz, timezone) if not type_def: @@ -314,7 +318,7 @@ class DateTime64(ChSqlaType, SqlaDateTime): SqlaDateTime.__init__(self) -class Time(ChSqlaType, Interval): +class Time(ChSqlaType, Interval): # type: ignore[misc] """ Represents the ClickHouse Time type, which corresponds to a timedelta. @@ -336,7 +340,7 @@ class Time(ChSqlaType, Interval): return None -class Time64(ChSqlaType, Interval): +class Time64(ChSqlaType, Interval): # type: ignore[misc] """ Represents the ClickHouse Time64 type with configurable precision. @@ -345,7 +349,7 @@ class Time64(ChSqlaType, Interval): If no precision is defined it default to 3. """ - def __init__(self, precision: int = None, type_def: TypeDef = None): + def __init__(self, precision: int | None = None, type_def: TypeDef | None = None): """ Time64 constructor with precision if not constructed with TypeDef. :param precision: 3 (ms), 6 (us), or 9 (ns) for sub-second precision. @@ -375,51 +379,29 @@ class Time64(ChSqlaType, Interval): return None -class Nullable: - """ - Class "wrapper" to use in DDL construction. It is never actually initialized but instead creates the "wrapped" - type with a Nullable wrapper - """ +def Nullable(element: ChSqlaType | type[ChSqlaType]) -> ChSqlaType: # noqa: N802 + """Wrap a ChSqlaType instance or class with a Nullable modifier for DDL construction.""" + if callable(element): + return element(type_def=NULLABLE_TYPE_DEF) + orig = element.type_def + wrappers = orig if "Nullable" in orig.wrappers else orig.wrappers + ("Nullable",) + return element.__class__(type_def=TypeDef(wrappers, orig.keys, orig.values)) - def __new__(cls, element: ChSqlaType | type[ChSqlaType]): - """ - Actually returns an instance of the enclosed type with a Nullable wrapper. If element is an instance, - constructs a new instance with a copied TypeDef plus the Nullable wrapper. If element is just a type, - constructs a new element of that type with only the Nullable wrapper. - :param element: ChSqlaType instance or class to wrap - """ - if callable(element): - return element(type_def=NULLABLE_TYPE_DEF) - orig = element.type_def - wrappers = orig if "Nullable" in orig.wrappers else orig.wrappers + ("Nullable",) - return element.__class__(type_def=TypeDef(wrappers, orig.keys, orig.values)) +def LowCardinality(element: ChSqlaType | type[ChSqlaType]) -> ChSqlaType: # noqa: N802 + """Wrap a ChSqlaType instance or class with a LowCardinality modifier for DDL construction.""" + if callable(element): + return element(type_def=LC_TYPE_DEF) + orig = element.type_def + wrappers = orig if "LowCardinality" in orig.wrappers else ("LowCardinality",) + orig.wrappers + return element.__class__(type_def=TypeDef(wrappers, orig.keys, orig.values)) -class LowCardinality: - """ - Class "wrapper" to use in DDL construction. It is never actually instantiated but instead creates the "wrapped" - type with a LowCardinality wrapper - """ - def __new__(cls, element: ChSqlaType | type[ChSqlaType]): - """ - Actually returns an instance of the enclosed type with a LowCardinality wrapper. If element is an instance, - constructs a new instance with a copied TypeDef plus the LowCardinality wrapper. If element is just a type, - constructs a new element of that type with only the LowCardinality wrapper. - :param element: ChSqlaType instance or class to wrap - """ - if callable(element): - return element(type_def=LC_TYPE_DEF) - orig = element.type_def - wrappers = orig if "LowCardinality" in orig.wrappers else ("LowCardinality",) + orig.wrappers - return element.__class__(type_def=TypeDef(wrappers, orig.keys, orig.values)) - - -class Array(ChSqlaType, ARRAY): +class Array(ChSqlaType, ARRAY): # type: ignore[misc] python_type = list dimensions = 1 - def __init__(self, element: ChSqlaType | type[ChSqlaType] = None, type_def: TypeDef = None): + def __init__(self, element: ChSqlaType | type[ChSqlaType] | None = None, type_def: TypeDef | None = None): """ Array constructor that can take a wrapped Array type if not constructed from a TypeDef :param element: ChSqlaType instance or class to wrap @@ -428,23 +410,25 @@ class Array(ChSqlaType, ARRAY): if not type_def: if callable(element): element = element() + if element is None: + raise ArgumentError("Array requires an element type or type_def") type_def = TypeDef(values=(element.name,)) ChSqlaType.__init__(self, type_def) # Set item_type directly; calling ARRAY.__init__ would reject nested Array(Array(T)), # which CH supports natively (CH expresses dimensions via nesting, not a dim count). # as_tuple has no class-level default, so set it here to satisfy ARRAY result processing. - self.item_type = sqla_type_from_name(type_def.values[0]) + self.item_type = cast("TypeEngine[Any]", sqla_type_from_name(type_def.values[0])) self.as_tuple = False -class Map(ChSqlaType, UserDefinedType): +class Map(ChSqlaType, UserDefinedType): # type: ignore[misc] python_type = dict def __init__( self, - key_type: ChSqlaType | type[ChSqlaType] = None, - value_type: ChSqlaType | type[ChSqlaType] = None, - type_def: TypeDef = None, + key_type: ChSqlaType | type[ChSqlaType] | None = None, + value_type: ChSqlaType | type[ChSqlaType] | None = None, + type_def: TypeDef | None = None, ): """ Map constructor that can take a wrapped key/values types if not constructed from a TypeDef @@ -457,18 +441,20 @@ class Map(ChSqlaType, UserDefinedType): key_type = key_type() if callable(value_type): value_type = value_type() + if key_type is None or value_type is None: + raise ArgumentError("Map requires key_type and value_type, or type_def") type_def = TypeDef(values=(key_type.name, value_type.name)) super().__init__(type_def) -class Tuple(ChSqlaType, UserDefinedType): +class Tuple(ChSqlaType, UserDefinedType): # type: ignore[misc] python_type = tuple def __init__( self, *args, - elements: Sequence[ChSqlaType | type[ChSqlaType]] = None, - type_def: TypeDef = None, + elements: Sequence[ChSqlaType | type[ChSqlaType]] | None = None, + type_def: TypeDef | None = None, ): """Tuple(UInt32, UUID) variadic form or Tuple(elements=[UInt32, UUID]) list form, not both.""" if type_def is None and not args and elements is None: @@ -481,7 +467,7 @@ class Tuple(ChSqlaType, UserDefinedType): raise ArgumentError("Cannot specify both positional elements and the 'elements' kwarg") if args: elements = args - values = [et() if callable(et) else et for et in elements] + values = [et() if callable(et) else et for et in elements] # type: ignore[union-attr] type_def = TypeDef(values=tuple(v.name for v in values)) super().__init__(type_def) @@ -493,7 +479,7 @@ class Tuple(ChSqlaType, UserDefinedType): return inst -class JSON(ChSqlaType, UserDefinedType): +class JSON(ChSqlaType, UserDefinedType): # type: ignore[misc] """ Note this isn't currently supported for insert/select, only table definitions """ @@ -501,7 +487,7 @@ class JSON(ChSqlaType, UserDefinedType): python_type = dict -class Nested(ChSqlaType, UserDefinedType): +class Nested(ChSqlaType, UserDefinedType): # type: ignore[misc] """ Note this isn't currently supported for insert/select, only table definitions """ @@ -509,14 +495,14 @@ class Nested(ChSqlaType, UserDefinedType): python_type = list -class SimpleAggregateFunction(ChSqlaType, UserDefinedType): +class SimpleAggregateFunction(ChSqlaType, UserDefinedType): # type: ignore[misc] python_type = str def __init__( self, - name: str = None, - element: ChSqlaType | type[ChSqlaType] = None, - type_def: TypeDef = None, + name: str | None = None, + element: ChSqlaType | type[ChSqlaType] | None = None, + type_def: TypeDef | None = None, ): """ Constructor that can take the SimpleAggregateFunction name and wrapped type if not constructed from a TypeDef @@ -527,23 +513,20 @@ class SimpleAggregateFunction(ChSqlaType, UserDefinedType): if not type_def: if callable(element): element = element() - type_def = TypeDef( - values=( - name, - element.name, - ) - ) + if element is None: + raise ArgumentError("SimpleAggregateFunction requires an element type or type_def") + type_def = TypeDef(values=(name, element.name)) super().__init__(type_def) -class AggregateFunction(ChSqlaType, UserDefinedType): +class AggregateFunction(ChSqlaType, UserDefinedType): # type: ignore[misc] """ Note this isn't currently supported for insert/select, only table definitions """ python_type = str - def __init__(self, *params, type_def: TypeDef = None): + def __init__(self, *params, type_def: TypeDef | None = None): """ Simply wraps the parameters for AggregateFunction for DDL, unless the TypeDef is specified. Callables or actual types are converted to their names. @@ -551,7 +534,7 @@ class AggregateFunction(ChSqlaType, UserDefinedType): :param type_def: TypeDef from parse_name function """ if not type_def: - values = () + values: tuple[Any, ...] = () for x in params: if callable(x): x = x() @@ -562,10 +545,10 @@ class AggregateFunction(ChSqlaType, UserDefinedType): super().__init__(type_def) -class QBit(ChSqlaType, UserDefinedType): +class QBit(ChSqlaType, UserDefinedType): # type: ignore[misc] python_type = list - def __init__(self, element_type: str = None, dimension: int = None, type_def: TypeDef = None): + def __init__(self, element_type: str | None = None, dimension: int | None = None, type_def: TypeDef | None = None): """ QBit constructor for bit-transposed vector types :param element_type: Element type (BFloat16, Float32, or Float64) diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/ddl/custom.py b/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/ddl/custom.py index b29e8e6b26a..2e7187ed6d6 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/ddl/custom.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/ddl/custom.py @@ -12,8 +12,8 @@ class CreateDatabase(DDL): def __init__( self, name: str, - engine: str = None, - zoo_path: str = None, + engine: str | None = None, + zoo_path: str | None = None, shard_name: str = "{shard}", replica_name: str = "{replica}", exists_ok: bool = False, diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/ddl/tableengine.py b/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/ddl/tableengine.py index 85759cfc6df..358c86aa0a5 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/ddl/tableengine.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/ddl/tableengine.py @@ -4,9 +4,8 @@ from typing import Any from sqlalchemy import Column from sqlalchemy.exc import ArgumentError, SQLAlchemyError -from sqlalchemy.sql.base import SchemaEventTarget from sqlalchemy.sql.elements import TextClause -from sqlalchemy.sql.visitors import Visitable +from sqlalchemy.sql.schema import SchemaItem from clickhouse_connect.cc_sqlalchemy.sql.sqlparse import split_top_level, walk_sql from clickhouse_connect.driver.binding import format_str, quote_identifier @@ -36,7 +35,7 @@ def _render_setting_value(value: Any) -> str: return format_str(str(value)) -def tuple_expr(expr_name, value: EngineParam): +def tuple_expr(expr_name: str, value: EngineParam) -> str: """ Create a table parameter with a tuple or list correctly formatted :param expr_name: parameter @@ -48,7 +47,7 @@ def tuple_expr(expr_name, value: EngineParam): v = f"{expr_name.strip()}" if isinstance(value, (tuple, list)): return f" {v} ({','.join(_render_engine_expr(item) for item in value)})" - return f"{v} {_render_engine_expr(value)}" + return f"{v} {_render_engine_expr(value)}" # type: ignore[arg-type] def repr_engine_value(value: Any) -> str: @@ -68,22 +67,22 @@ def repr_engine_value(value: Any) -> str: return repr(value) -class TableEngine(SchemaEventTarget, Visitable): +class TableEngine(SchemaItem): """ SqlAlchemy Schema element to support ClickHouse table engines. At the moment provides no real functionality other than the CREATE TABLE argument string """ - arg_names = () - quoted_args = set() - optional_args = set() - eng_params = () + arg_names: Sequence[str] = () + quoted_args: set[str] = set() + optional_args: set[str] = set() + eng_params: Sequence[str] = () def __init_subclass__(cls, **kwargs): engine_map[cls.__name__] = cls def __init__(self, kwargs): - Visitable.__init__(self) + super().__init__() self.name = self.__class__.__name__ te_name = f"{self.name} Table Engine" self._orig_kwargs = kwargs.copy() @@ -167,21 +166,21 @@ class Set(TableEngine): class Dictionary(TableEngine): arg_names = ["dictionary"] - def __init__(self, dictionary: str = None): + def __init__(self, dictionary: str | None = None): super().__init__(locals()) class Merge(TableEngine): arg_names = ["db_name, tables_regexp"] - def __init__(self, db_name: str = None, tables_regexp: str = None): + def __init__(self, db_name: str | None = None, tables_regexp: str | None = None): super().__init__(locals()) class File(TableEngine): arg_names = ["fmt"] - def __init__(self, fmt: str = None): + def __init__(self, fmt: str | None = None): super().__init__(locals()) @@ -189,7 +188,14 @@ class Distributed(TableEngine): arg_names = ["cluster", "database", "table", "sharding_key", "policy_name"] optional_args = {"sharding_key", "policy_name"} - def __init__(self, cluster: str = None, database: str = None, table=None, sharding_key: str = None, policy_name: str = None): + def __init__( + self, + cluster: str | None = None, + database: str | None = None, + table=None, + sharding_key: str | None = None, + policy_name: str | None = None, + ): super().__init__(locals()) @@ -229,9 +235,9 @@ class ReplacingMergeTree(TableEngine): def __init__( self, - ver: str = None, - version: str = None, - is_deleted: str = None, + ver: str | None = None, + version: str | None = None, + is_deleted: str | None = None, order_by: EngineParam = None, primary_key: EngineParam = None, partition_by: EngineParam = None, @@ -260,7 +266,7 @@ class CollapsingMergeTree(TableEngine): def __init__( self, - sign: str = None, + sign: str | None = None, order_by: EngineParam = None, primary_key: EngineParam = None, partition_by: EngineParam = None, @@ -279,8 +285,8 @@ class VersionedCollapsingMergeTree(TableEngine): def __init__( self, - sign: str = None, - version: str = None, + sign: str | None = None, + version: str | None = None, order_by: EngineParam = None, primary_key: EngineParam = None, partition_by: EngineParam = None, @@ -300,8 +306,8 @@ class GraphiteMergeTree(TableEngine): def __init__( self, - config_section: str = None, - version: str = None, + config_section: str | None = None, + version: str | None = None, order_by: EngineParam = None, primary_key: EngineParam = None, partition_by: EngineParam = None, @@ -326,8 +332,8 @@ class ReplicatedMergeTree(TableEngine): primary_key: EngineParam = None, partition_by: EngineParam = None, sample_by: EngineParam = None, - zk_path: str = None, - replica: str = None, + zk_path: str | None = None, + replica: str | None = None, ttl: EngineExpr | None = None, settings: dict[str, Any] | None = None, ): @@ -352,13 +358,13 @@ class ReplicatedReplacingMergeTree(TableEngine): def __init__( self, - ver: str = None, + ver: str | None = None, order_by: EngineParam = None, primary_key: EngineParam = None, partition_by: EngineParam = None, sample_by: EngineParam = None, - zk_path: str = None, - replica: str = None, + zk_path: str | None = None, + replica: str | None = None, ttl: EngineExpr | None = None, settings: dict[str, Any] | None = None, ): @@ -375,13 +381,13 @@ class ReplicatedCollapsingMergeTree(TableEngine): def __init__( self, - sign: str = None, + sign: str | None = None, order_by: EngineParam = None, primary_key: EngineParam = None, partition_by: EngineParam = None, sample_by: EngineParam = None, - zk_path: str = None, - replica: str = None, + zk_path: str | None = None, + replica: str | None = None, ttl: EngineExpr | None = None, settings: dict[str, Any] | None = None, ): @@ -398,14 +404,14 @@ class ReplicatedVersionedCollapsingMergeTree(TableEngine): def __init__( self, - sign: str = None, - version: str = None, + sign: str | None = None, + version: str | None = None, order_by: EngineParam = None, primary_key: EngineParam = None, partition_by: EngineParam = None, sample_by: EngineParam = None, - zk_path: str = None, - replica: str = None, + zk_path: str | None = None, + replica: str | None = None, ttl: EngineExpr | None = None, settings: dict[str, Any] | None = None, ): @@ -422,13 +428,13 @@ class ReplicatedGraphiteMergeTree(TableEngine): def __init__( self, - config_section: str = None, + config_section: str | None = None, order_by: EngineParam = None, primary_key: EngineParam = None, partition_by: EngineParam = None, sample_by: EngineParam = None, - zk_path: str = None, - replica: str = None, + zk_path: str | None = None, + replica: str | None = None, ttl: EngineExpr | None = None, settings: dict[str, Any] | None = None, ): diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/dialect.py b/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/dialect.py index 368e80422ba..bdfd529af8d 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/dialect.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/dialect.py @@ -31,7 +31,7 @@ class ClickHouseDialect(DefaultDialect): postfetch_lastrowid = False ddl_compiler = ChDDLCompiler statement_compiler = ChStatementCompiler - preparer = ChIdentifierPreparer + preparer = ChIdentifierPreparer # type: ignore[assignment] description_encoding = None max_identifier_length = 127 ischema_names = ischema_names diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/sql/__init__.py b/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/sql/__init__.py index ec2cee4c942..3352936603f 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/sql/__init__.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/sql/__init__.py @@ -1,3 +1,5 @@ +from typing import Any + from sqlalchemy import Table, and_ from sqlalchemy.sql.selectable import FromClause, Select @@ -16,7 +18,7 @@ def full_table(table_name: str, schema: str | None = None) -> str: return f"{quote_identifier(schema)}.{quote_identifier(table_name)}" -def format_table(table: Table): +def format_table(table: Table) -> str: return full_table(table.name, table.schema) @@ -45,21 +47,21 @@ def _resolve_target(select_stmt: Select, table: FromClause | None, method_name: def _target_cache_key(target: FromClause) -> str: if hasattr(target, "fullname"): - return target.fullname - return target.name + return target.fullname # type: ignore[attr-defined] + return target.name # type: ignore[attr-defined] def final(select_stmt: Select, table: FromClause | None = None) -> Select: """Apply the ClickHouse FINAL modifier. For ReplacingMergeTree-family engines.""" target = _resolve_target(select_stmt, table, "final") - ch_final = getattr(select_stmt, "_ch_final", set()) + ch_final: set[FromClause] = getattr(select_stmt, "_ch_final", set()) if target in ch_final: return select_stmt hint_key = _target_cache_key(target) new_stmt = select_stmt.with_statement_hint(f"FINAL:{hint_key}", dialect_name=_CH_MODIFIER_DIALECT) - new_stmt._ch_final = ch_final | {target} + new_stmt._ch_final = ch_final | {target} # type: ignore[attr-defined] return new_stmt @@ -75,7 +77,7 @@ def sample(select_stmt: Select, sample_value: str | int | float, table: FromClau new_stmt = select_stmt.with_statement_hint(f"SAMPLE:{hint_key}:{sample_value}", dialect_name=_CH_MODIFIER_DIALECT) ch_sample = dict(getattr(select_stmt, "_ch_sample", {})) ch_sample[target] = sample_value - new_stmt._ch_sample = ch_sample + new_stmt._ch_sample = ch_sample # type: ignore[attr-defined] return new_stmt @@ -83,7 +85,7 @@ def _select_sample(self: Select, sample_value: str | int | float, table: FromCla return sample(self, sample_value=sample_value, table=table) -def _apply_array_join(select_stmt: Select, cols, alias, is_left: bool) -> Select: +def _apply_array_join(select_stmt: Select, cols: Any, alias: Any, is_left: bool) -> Select: if not isinstance(select_stmt, Select): raise TypeError("array_join() expects a SQLAlchemy Select instance") @@ -125,7 +127,7 @@ def _select_left_array_join(self: Select, *cols, alias=None) -> Select: return _apply_array_join(self, cols, alias, is_left=True) -def prewhere(select_stmt, whereclause): +def prewhere(select_stmt: Select, whereclause: Any) -> Select: """Apply ClickHouse PREWHERE. Multiple calls compose with AND.""" if not isinstance(select_stmt, Select): raise TypeError("prewhere() expects a SQLAlchemy Select instance") @@ -136,11 +138,11 @@ def prewhere(select_stmt, whereclause): # Hint key is str(combined) (structural, with bind placeholders) rather # than id() so equivalent statements share a compiled-statement cache entry. new_stmt = select_stmt.with_statement_hint(f"PREWHERE:{str(combined)}", dialect_name=_CH_MODIFIER_DIALECT) - new_stmt._ch_prewhere = PreWhereClause(combined) + new_stmt._ch_prewhere = PreWhereClause(combined) # type: ignore[attr-defined] return new_stmt -def limit_by(select_stmt, by_clauses, limit, offset=None): +def limit_by(select_stmt: Select, by_clauses: Any, limit: int, offset: int | None = None) -> Select: """Apply ClickHouse LIMIT BY (top-N per group). Renders `LIMIT [offset,] limit BY by_clauses`.""" if not isinstance(select_stmt, Select): raise TypeError("limit_by() expects a SQLAlchemy Select instance") @@ -151,21 +153,21 @@ def limit_by(select_stmt, by_clauses, limit, offset=None): by_key = ",".join(str(c) for c in by_tuple) new_stmt = select_stmt.with_statement_hint(f"LIMIT_BY:{limit}:{offset}:{by_key}", dialect_name=_CH_MODIFIER_DIALECT) - new_stmt._ch_limit_by = LimitByClause(by_tuple, limit, offset) + new_stmt._ch_limit_by = LimitByClause(by_tuple, limit, offset) # type: ignore[attr-defined] return new_stmt -def _select_prewhere(self, whereclause): +def _select_prewhere(self: Select, whereclause: Any) -> Select: return prewhere(self, whereclause) -def _select_limit_by(self, by_clauses, limit, offset=None): +def _select_limit_by(self: Select, by_clauses: Any, limit: int, offset: int | None = None) -> Select: return limit_by(self, by_clauses, limit, offset) -Select.sample = _select_sample -Select.final = _select_final -Select.array_join = _select_array_join -Select.left_array_join = _select_left_array_join -Select.prewhere = _select_prewhere -Select.limit_by = _select_limit_by +Select.sample = _select_sample # type: ignore[attr-defined] +Select.final = _select_final # type: ignore[attr-defined] +Select.array_join = _select_array_join # type: ignore[attr-defined] +Select.left_array_join = _select_left_array_join # type: ignore[attr-defined] +Select.prewhere = _select_prewhere # type: ignore[attr-defined] +Select.limit_by = _select_limit_by # type: ignore[attr-defined] diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/sql/preparer.py b/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/sql/preparer.py index 4521eb97dfa..cda0f5a2ed8 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/sql/preparer.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/sql/preparer.py @@ -4,7 +4,7 @@ from clickhouse_connect.driver.binding import quote_identifier class ChIdentifierPreparer(IdentifierPreparer): - quote_identifier = staticmethod(quote_identifier) + quote_identifier = staticmethod(quote_identifier) # type: ignore[assignment] def __init__(self, dialect, **kwargs): super().__init__(dialect, **kwargs) diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/datatypes/base.py b/contrib/python/clickhouse-connect/clickhouse_connect/datatypes/base.py index 220a2dda579..ab5e17740c4 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/datatypes/base.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/datatypes/base.py @@ -16,8 +16,8 @@ from clickhouse_connect.driver.query import QueryContext from clickhouse_connect.driver.types import ByteSource logger = logging.getLogger(__name__) -ch_read_formats = {} -ch_write_formats = {} +ch_read_formats: dict[type, str] = {} +ch_write_formats: dict[type, str] = {} class TypeDef(NamedTuple): @@ -45,10 +45,10 @@ class ClickHouseType(ABC): # noqa: B024 np_type = "O" # Default to Numpy Object type nano_divisor = 0 # Only relevant for date like objects byte_size = 0 - valid_formats = "native" + valid_formats: str | tuple[str, ...] = "native" - python_type = None - base_type = None + python_type: type | None = None + base_type: str | None = None @property def _null_time_unit(self): @@ -110,7 +110,7 @@ class ClickHouseType(ABC): # noqa: B024 def insert_name(self): return self.name - def data_size(self, sample: Sequence) -> int: + def data_size(self, sample: Collection) -> int: if self.low_card: values = set(sample) d_size = self._data_size(values) + 2 @@ -126,7 +126,7 @@ class ClickHouseType(ABC): # noqa: B024 total = 0 for x in sample: total += len(str(x)) - return total / len(sample) + 1 + return total // len(sample) + 1 def write_column_prefix(self, dest: bytearray): """ @@ -264,9 +264,9 @@ class ClickHouseType(ABC): # noqa: B024 def _write_column_low_card(self, column: Sequence, dest: bytearray, ctx: InsertContext): if len(column) == 0: return - keys = [] - index = [] - rev_map = {} + keys: list[int] = [] + index: list[Any] = [] + rev_map: dict[Any, int] = {} rmg = rev_map.get if self.nullable: index.append(None) @@ -319,15 +319,16 @@ class ArrayType(ClickHouseType, ABC, registered=False): """ _signed = True - _array_type = None - _struct_type = None + _array_type: str | None = None + _struct_type: str | None = None valid_formats = "string", "native" - python_type = int + python_type: type = int def __init_subclass__(cls, registered: bool = True): super().__init_subclass__(registered) if cls._array_type in ("i", "I") and int_size == 2: - cls._array_type = "L" if cls._array_type.isupper() else "l" + array_type_char = cls._array_type + cls._array_type = "L" if array_type_char.isupper() else "l" if isinstance(cls._array_type, str) and cls._array_type: cls._struct_type = "<" + cls._array_type cls.byte_size = array.array(cls._array_type).itemsize @@ -335,14 +336,17 @@ class ArrayType(ClickHouseType, ABC, registered=False): def _read_column_binary(self, source: ByteSource, num_rows: int, ctx: QueryContext, _read_state: Any): if ctx.use_numpy: return driver_ctypes.numpy_conv.read_numpy_array(source, self.np_type, num_rows) + assert self._array_type is not None return source.read_array(self._array_type, num_rows) def _read_nullable_column(self, source: ByteSource, num_rows: int, ctx: QueryContext, _read_state: Any) -> Sequence: + assert self._array_type is not None return data_conv.read_nullable_array(source, self._array_type, num_rows, self._active_null(ctx)) def _build_lc_column(self, index: Sequence, keys: array.array, ctx: QueryContext): if ctx.use_numpy: - return options.np.fromiter((index[key] for key in keys), dtype=index.dtype, count=len(index)) + # index is a numpy array when ctx.use_numpy is True + return options.np.fromiter((index[key] for key in keys), dtype=index.dtype, count=len(index)) # type: ignore[attr-defined] return super()._build_lc_column(index, keys, ctx) def _finalize_column(self, column: Sequence, ctx: QueryContext) -> Sequence: @@ -357,6 +361,7 @@ class ArrayType(ClickHouseType, ABC, registered=False): def _write_column_binary(self, column: Sequence | MutableSequence, dest: bytearray, ctx: InsertContext): if len(column) and self.nullable: column = [0 if x is None else x for x in column] + assert self._array_type is not None write_array(self._array_type, column, dest, ctx.column_name) def _active_null(self, ctx: QueryContext): @@ -377,7 +382,7 @@ class UnsupportedType(ClickHouseType, ABC, registered=False): super().__init__(type_def) self._name_suffix = type_def.arg_str - def _read_column_binary(self, source: Sequence, num_rows: int, ctx: QueryContext, read_state: Any): + def _read_column_binary(self, _source: ByteSource, _num_rows: int, _ctx: QueryContext, _read_state: Any): raise NotSupportedError(f"{self.name} deserialization not supported") def _write_column_binary(self, column: Sequence | MutableSequence, dest: bytearray, ctx: InsertContext): diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/datatypes/container.py b/contrib/python/clickhouse-connect/clickhouse_connect/datatypes/container.py index 88e87540695..951e618d623 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/datatypes/container.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/datatypes/container.py @@ -33,7 +33,7 @@ class Array(ClickHouseType): def read_column_prefix(self, source: ByteSource, ctx: QueryContext): return self.element_type.read_column_prefix(source, ctx) - def _data_size(self, sample: Sequence) -> int: + def _data_size(self, sample: Collection[Any]) -> int: if len(sample) == 0: return 8 total = 0 @@ -139,7 +139,7 @@ class Tuple(ClickHouseType): column = e_type.read_column_data(source, num_rows, ctx, read_state[ix]) columns.append(column) if e_names and self.read_format(ctx) != "tuple": - dicts = [{} for _ in range(num_rows)] + dicts: list[dict[str, Any]] = [{} for _ in range(num_rows)] for ix, x in enumerate(dicts): for y, key in enumerate(e_names): x[key] = columns[y][ix] @@ -163,7 +163,7 @@ class Tuple(ClickHouseType): def convert_dict_insert(self, column: Sequence) -> Sequence: names = self.element_names - col = [[] for _ in names] + col: list[list[Any]] = [[] for _ in names] for x in column: for ix, name in enumerate(names): col[ix].append(x.get(name)) diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/datatypes/dynamic.py b/contrib/python/clickhouse-connect/clickhouse_connect/datatypes/dynamic.py index 150fa1ced1f..4732ff1b957 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/datatypes/dynamic.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/datatypes/dynamic.py @@ -156,7 +156,7 @@ class Variant(ClickHouseType): v_count = len(self.element_types) if v_count == 0: return 1 - sub_samples = [[] for _ in range(v_count)] + sub_samples: list[list[Any]] = [[] for _ in range(v_count)] for v in sample: if v is None: continue @@ -197,7 +197,7 @@ def read_variant_column( # Now we have to walk through each of the discriminators again to assign the correct value from # the sub-column to the final result column sub_indexes = [0] * v_count - col = [] + col: list[Any] = [] app_col = col.append for disc in discriminators: if disc == 255: @@ -230,7 +230,9 @@ def read_dynamic_prefix(_, source: ByteSource, ctx: QueryContext) -> DynamicStat class Dynamic(ClickHouseType): python_type = object - read_column_prefix = read_dynamic_prefix + + def read_column_prefix(self, source: ByteSource, ctx: QueryContext) -> DynamicState: + return read_dynamic_prefix(self, source, ctx) @property def insert_name(self): @@ -272,7 +274,7 @@ def write_json(ch_type: ClickHouseType, column: Sequence, dest: bytearray, ctx: first = first_value(column, ch_type.nullable) write_col = column - encoding = ctx.encoding or ch_type.encoding + encoding: str | None = ctx.encoding or ch_type.encoding if not isinstance(first, str) and ch_type.write_format(ctx) != "string": to_json = any_to_json if ch_type.nullable: @@ -576,7 +578,7 @@ class JSON(ClickHouseType): shared_columns = SHARED_DATA_TYPE.read_column_data(source, num_rows, ctx, read_state.shared_state) # noqa: F821 (undefined-name) col = [] for row_num in range(num_rows): - top = {} + top: dict[str, Any] = {} for ix, field in enumerate(self.typed_paths): _nest_value(top, field, typed_columns[ix][row_num]) for ix, field in enumerate(read_state.dynamic_paths): diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/datatypes/format.py b/contrib/python/clickhouse-connect/clickhouse_connect/datatypes/format.py index fa09e2c9fdd..9bb02e72b7d 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/datatypes/format.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/datatypes/format.py @@ -63,7 +63,7 @@ def _convert_arguments(*args, **kwargs) -> dict[str, str]: return fmt_map -def _matching_types(pattern: str, fmt: str = None) -> Sequence[type[ClickHouseType]]: +def _matching_types(pattern: str, fmt: str | None = None) -> Sequence[type[ClickHouseType]]: re_pattern = re.compile(pattern.replace("*", ".*"), re.IGNORECASE) matches = [ch_type for type_name, ch_type in type_map.items() if re_pattern.match(type_name)] if not matches: diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/datatypes/network.py b/contrib/python/clickhouse-connect/clickhouse_connect/datatypes/network.py index e0e751cf390..69f200dd7b9 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/datatypes/network.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/datatypes/network.py @@ -67,22 +67,24 @@ class IPv6(ClickHouseType): """Read IPv6 addresses in native format, always returning IPv6Address objects.""" fast_ip_v6 = IPv6Address.__new__ with_scope_id = "_scope_id" in IPv6Address.__slots__ - new_col = [] + new_col: list[IPv6Address] = [] app = new_col.append ifb = int.from_bytes for _ in range(num_rows): int_value = ifb(source.read_bytes(16), "big") ipv6 = fast_ip_v6(IPv6Address) - ipv6._ip = int_value + # Bypass IPv6Address.__init__ for performance; _ip and _scope_id are + # the internal representation used by CPython's ipaddress module. + ipv6._ip = int_value # type: ignore[attr-defined] if with_scope_id: - ipv6._scope_id = None + ipv6._scope_id = None # type: ignore[attr-defined] app(ipv6) return new_col @staticmethod def _read_binary_str(source: ByteSource, num_rows: int) -> list[str]: """Read IPv6 addresses in string format, always returning IPv6Address strings.""" - new_col = [] + new_col: list[str] = [] app = new_col.append tov6 = socket.inet_ntop af6 = socket.AF_INET6 diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/datatypes/numeric.py b/contrib/python/clickhouse-connect/clickhouse_connect/datatypes/numeric.py index 37b0d83ed4c..90d2abea7ab 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/datatypes/numeric.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/datatypes/numeric.py @@ -16,6 +16,8 @@ from clickhouse_connect.driver.types import ByteSource class IntBase(ArrayType, registered=False): + _array_type: str + def _write_column_binary(self, column: Sequence | MutableSequence, dest: bytearray, ctx: InsertContext): if len(column) == 0: return @@ -109,7 +111,7 @@ class BigInt(ClickHouseType, registered=False): def _read_column_binary(self, source: ByteSource, num_rows: int, ctx: QueryContext, _read_state: Any): signed = self._signed sz = self.byte_size - column = [] + column: list[Any] = [] app = column.append ifb = int.from_bytes if self.read_format(ctx) == "string": @@ -407,7 +409,7 @@ class BigDecimal(Decimal, registered=False): def _read_column_binary(self, source: ByteSource, num_rows: int, _ctx: QueryContext, _read_state: Any): dec = decimal.Decimal scale = self.scale - column = [] + column: list[Any] = [] app = column.append sz = self.byte_size ifb = int.from_bytes diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/datatypes/registry.py b/contrib/python/clickhouse-connect/clickhouse_connect/datatypes/registry.py index 3f5e45866c8..cb035782c85 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/datatypes/registry.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/datatypes/registry.py @@ -1,4 +1,5 @@ import logging +from typing import Any from clickhouse_connect.datatypes.base import ClickHouseType, TypeDef, type_map from clickhouse_connect.driver.exceptions import InternalError @@ -18,7 +19,8 @@ def parse_name(name: str) -> tuple[str, str, TypeDef]: """ base = name wrappers = [] - keys = tuple() + keys: tuple[Any, ...] = () + values: tuple[Any, ...] = () if base.startswith("LowCardinality"): wrappers.append("LowCardinality") base = base[15:-1] diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/datatypes/special.py b/contrib/python/clickhouse-connect/clickhouse_connect/datatypes/special.py index 7b030fb3662..f5f5df9f45d 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/datatypes/special.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/datatypes/special.py @@ -1,4 +1,4 @@ -from collections.abc import MutableSequence, Sequence +from collections.abc import Collection, MutableSequence, Sequence from typing import Any from uuid import UUID as PYUUID @@ -30,7 +30,7 @@ class UUID(ClickHouseType): @staticmethod def _read_binary_str(source: ByteSource, num_rows: int): v = source.read_array("Q", num_rows * 2) - column = [] + column: list[str] = [] app = column.append for i in range(num_rows): ix = i << 1 @@ -94,7 +94,7 @@ class SimpleAggregateFunction(ClickHouseType): self.python_type = self.element_type.python_type self.nano_divisor = self.element_type.nano_divisor - def _data_size(self, sample: Sequence) -> int: + def _data_size(self, sample: Collection[Any]) -> int: return self.element_type.data_size(sample) def read_column_prefix(self, source: ByteSource, ctx: QueryContext): diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/datatypes/temporal.py b/contrib/python/clickhouse-connect/clickhouse_connect/datatypes/temporal.py index 0f8ed6ad43f..c2105643afb 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/datatypes/temporal.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/datatypes/temporal.py @@ -4,9 +4,9 @@ import array import re import zoneinfo from abc import abstractmethod -from collections.abc import MutableSequence, Sequence +from collections.abc import Callable, MutableSequence, Sequence from datetime import date, datetime, time, timedelta, tzinfo -from typing import TYPE_CHECKING, Any, NamedTuple +from typing import TYPE_CHECKING, Any, NamedTuple, cast if TYPE_CHECKING: import numpy @@ -50,6 +50,7 @@ class Date(ClickHouseType): if self.nullable: column = [x if x else 0 for x in column] else: + esd: date if isinstance(first, datetime): esd = epoch_start_datetime else: @@ -112,6 +113,7 @@ class Date32(Date): class DateTimeBase(ClickHouseType, registered=False): __slots__ = ("tzinfo",) + tzinfo: tzinfo | None valid_formats = "native", "int" python_type = datetime @@ -143,12 +145,13 @@ class DateTimeBase(ClickHouseType, registered=False): if isinstance(column, list): column = options.pd.DatetimeIndex(column) - if column.tz is None: - result = column.astype(self.pandas_dtype) + dti = cast(Any, column) + if dti.tz is None: + result = dti.astype(self.pandas_dtype) return options.pd.array(result) if self.nullable else result - naive_ns = column.tz_convert("UTC").tz_localize(None).astype(self.pandas_dtype) - tz_aware_result = naive_ns.tz_localize("UTC").tz_convert(column.tz) + naive_ns = dti.tz_convert("UTC").tz_localize(None).astype(self.pandas_dtype) + tz_aware_result = naive_ns.tz_localize("UTC").tz_convert(dti.tz) return options.pd.array(tz_aware_result) if self.nullable else tz_aware_result if self.nullable: @@ -329,7 +332,6 @@ class TimeBase(ClickHouseType, registered=False): _array_type: str byte_size: int - np_type: str valid_formats = ("native", "string", "int", "time") python_type = timedelta @@ -403,7 +405,7 @@ class TimeBase(ClickHouseType, registered=False): return [0] * len(column) return [] - converter_map = { + converter_map: dict[type, Callable[..., int]] = { timedelta: self._timedelta_to_ticks, time: self._time_to_ticks, float: self._numerical_to_ticks, @@ -509,7 +511,7 @@ class TimeBase(ClickHouseType, registered=False): raise NotImplementedError @abstractmethod - def _ticks_to_np_timedelta(self, ticks: int) -> timedelta: + def _ticks_to_np_timedelta(self, ticks: int) -> timedelta | numpy.timedelta64: """Convert integer ticks into an np.timedelta.""" raise NotImplementedError @@ -638,7 +640,7 @@ class Time64(TimeBase): return self._SECONDS_PER_DAY * self.precision - 1 @property - def np_type(self) -> str: + def np_type(self): return f"timedelta64{self.unit}" @property diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/datatypes/vector.py b/contrib/python/clickhouse-connect/clickhouse_connect/datatypes/vector.py index 0d666941f74..e0dff1b56c7 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/datatypes/vector.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/datatypes/vector.py @@ -4,7 +4,7 @@ import logging from collections.abc import Sequence from math import ceil, nan from struct import pack, unpack -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast if TYPE_CHECKING: import numpy @@ -89,7 +89,7 @@ class QBit(ClickHouseType): tuple_data = self._tuple_type.read_column_data(source, num_rows, ctx, read_state) vectors = [self._untranspose_row(t) for t in tuple_data] if self.nullable: - return data_conv.build_nullable_column(vectors, null_map, self._active_null(ctx)) + return data_conv.build_nullable_column(vectors, cast(bytes, null_map), self._active_null(ctx)) return vectors def write_column_prefix(self, dest: bytearray): diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/dbapi/__init__.py b/contrib/python/clickhouse-connect/clickhouse_connect/dbapi/__init__.py index 84a9a8c3fcc..99a298d4f0d 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/dbapi/__init__.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/dbapi/__init__.py @@ -1,3 +1,5 @@ +from typing import Any + from clickhouse_connect.dbapi.connection import Connection apilevel = "2.0" # PEP 249 DB API level @@ -12,11 +14,11 @@ class Error(Exception): def connect( host: str | None = None, database: str | None = None, - username: str | None = "", - password: str | None = "", + username: str = "", + password: str = "", port: int | None = None, - **kwargs, -): + **kwargs: Any, +) -> Connection: secure = kwargs.pop("secure", False) return Connection( host=host, diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/dbapi/connection.py b/contrib/python/clickhouse-connect/clickhouse_connect/dbapi/connection.py index 3dba3c54aef..7b0cafaf121 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/dbapi/connection.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/dbapi/connection.py @@ -1,3 +1,5 @@ +from typing import Any + from clickhouse_connect.dbapi.cursor import Cursor from clickhouse_connect.driver import create_client from clickhouse_connect.driver.query import QueryResult @@ -10,15 +12,15 @@ class Connection: def __init__( self, - dsn: str = None, + dsn: str | None = None, username: str = "", password: str = "", - host: str = None, - database: str = None, - interface: str = None, - port: int = 0, + host: str | None = None, + database: str | None = None, + interface: str | None = None, + port: int | None = None, secure: bool | str = False, - **kwargs, + **kwargs: Any, ): self.client = create_client( host=host, @@ -35,20 +37,20 @@ class Connection: self.client._add_integration_tag("sqlalchemy") self.timezone = self.client.server_tz - def close(self): + def close(self) -> None: self.client.close() - def commit(self): + def commit(self) -> None: pass - def rollback(self): + def rollback(self) -> None: pass - def command(self, cmd: str): + def command(self, cmd: str) -> Any: return self.client.command(cmd) def raw_query(self, query: str) -> QueryResult: return self.client.query(query) - def cursor(self): + def cursor(self) -> Cursor: return Cursor(self.client) diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/dbapi/cursor.py b/contrib/python/clickhouse-connect/clickhouse_connect/dbapi/cursor.py index af7f36f67b6..9b200aba3a3 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/dbapi/cursor.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/dbapi/cursor.py @@ -1,6 +1,7 @@ import logging import re from collections.abc import Mapping, Sequence +from typing import Any, cast from clickhouse_connect.datatypes.registry import get_from_name from clickhouse_connect.driver import Client @@ -23,34 +24,34 @@ class Cursor: def __init__(self, client: Client): self.client = client - self.arraysize = 1 + self.arraysize: int = 1 self.data: Sequence | None = None - self.names = [] - self.types = [] - self._rowcount = 0 - self._summary: list[dict[str, str]] = [] - self._ix = 0 + self.names: Sequence[str] = [] + self.types: Sequence[Any] = [] + self._rowcount: int = 0 + self._summary: list[dict[str, Any]] = [] + self._ix: int = 0 - def check_valid(self): + def check_valid(self) -> None: if self.data is None: raise ProgrammingError("Cursor is not valid") @property - def description(self): + def description(self) -> list[tuple[str, Any, None, None, None, None, bool]]: return [(n, t, None, None, None, None, True) for n, t in zip(self.names, self.types)] @property - def rowcount(self): + def rowcount(self) -> int: return self._rowcount @property - def summary(self) -> list[dict[str, str]]: + def summary(self) -> list[dict[str, Any]]: return self._summary - def close(self): + def close(self) -> None: self.data = None - def execute(self, operation: str, parameters=None): + def execute(self, operation: str, parameters: Any = None) -> None: if not parameters and isinstance(operation, str): # Per PEP 249 pyformat paramstyle, callers (e.g. SQLAlchemy) escape # literal percent signs as %% in operation strings. When there are @@ -80,7 +81,7 @@ class Cursor: self.names = meta_result.column_names self.types = [x.name for x in meta_result.column_types] - def _try_bulk_insert(self, operation: str, data): + def _try_bulk_insert(self, operation: str, data: Any) -> bool: match = insert_re.match(remove_sql_comments(operation)) if not match: return False @@ -94,24 +95,31 @@ class Cursor: op_columns = None if "VALUES" not in temp.upper(): return False + if not isinstance(data, Sequence) or len(data) == 0: + return False first_row = data[0] + col_names: list[str] | str + data_values: Sequence[Sequence[Any]] if isinstance(first_row, Mapping): - col_names = list(first_row.keys()) - if op_columns and {unescape_identifier(x) for x in op_columns} != set(col_names): + col_names = [str(k) for k in first_row.keys()] + if op_columns and {unescape_identifier(str(x)) for x in op_columns} != set(col_names): return False # Data sent in doesn't match the columns in the insert statement data_values = [list(row.values()) for row in data] elif isinstance(first_row, Sequence) and not isinstance(first_row, (str, bytes)): # PEP 249 also allows rows as sequences; take column names from the # insert statement if present, otherwise insert into all columns - col_names = [unescape_identifier(x) for x in op_columns] if op_columns else "*" + col_names = [unescape_identifier(str(x)) for x in op_columns] if op_columns else "*" data_values = data else: return False - self.client.insert(table, data_values, col_names) + insert_summary = self.client.insert(table, data_values, col_names) self.data = [] + self._rowcount = insert_summary.written_rows + self._ix = 0 + self._summary.append(insert_summary.summary) return True - def executemany(self, operation, parameters): + def executemany(self, operation: str, parameters: Any) -> None: if not parameters or self._try_bulk_insert(operation, parameters): return self.data = [] @@ -138,22 +146,25 @@ class Cursor: # Need to reset cursor _ix after performing an execute self._ix = 0 - def fetchall(self): + def fetchall(self) -> Sequence: self.check_valid() - ret = self.data[self._ix :] + data = cast(Sequence, self.data) + ret = data[self._ix :] self._ix = self._rowcount return ret - def fetchone(self): + def fetchone(self) -> Any: self.check_valid() if self._ix >= self._rowcount: return None - val = self.data[self._ix] + data = cast(Sequence, self.data) + val = data[self._ix] self._ix += 1 return val - def fetchmany(self, size: int = -1): + def fetchmany(self, size: int = -1) -> Sequence: self.check_valid() + data = cast(Sequence, self.data) if size < 0: # Fetch all remaining rows @@ -163,12 +174,12 @@ class Cursor: return [] end = min(self._ix + size, self._rowcount) - ret = self.data[self._ix : end] + ret = data[self._ix : end] self._ix = end return ret - def nextset(self): + def nextset(self) -> None: raise NotImplementedError - def callproc(self, *args, **kwargs): + def callproc(self, *args, **kwargs) -> None: raise NotImplementedError diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/driver/__init__.py b/contrib/python/clickhouse-connect/clickhouse_connect/driver/__init__.py index 811286ad24c..1775a420230 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/driver/__init__.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/driver/__init__.py @@ -44,23 +44,24 @@ def _parse_connection_params( host: str | None, username: str | None, password: str, - port: int, - database: str, + port: int | None, + database: str | None, interface: str | None, secure: bool | str, dsn: str | None, kwargs: dict[str, Any], -) -> tuple[str, str | None, str, int, str, str]: +) -> tuple[str, str | None, str, int, str | None, str]: """Parse and normalize connection parameters including DSN parsing.""" + if database == "__default__": # legacy sentinel for "not specified" + database = None if dsn: parsed = urlparse(dsn) username = username or _unquote(parsed.username) password = password or _unquote(parsed.password) or "" host = host or parsed.hostname port = port or parsed.port - if parsed.path and (not database or database == "__default__"): - database = unquote(parsed.path[1:].split("/")[0]) - database = database or parsed.path + if not database and parsed.path: + database = unquote(parsed.path[1:].split("/")[0]) or None for k, v in parse_qs(parsed.query).items(): kwargs[k] = v[0] use_tls = str(secure).lower() == "true" or interface == "https" or (not interface and str(port) in ("443", "8443")) @@ -81,7 +82,12 @@ def _parse_connection_params( return host, username, password, port, database, interface -def _validate_access_token(access_token: str | None, token_provider: Callable[[], str] | None, username: str | None, password: str) -> None: +def _validate_access_token( + access_token: str | None, + token_provider: Callable[[], str | Awaitable[str]] | None, + username: str | None, + password: str, +) -> None: """Validate that token-based and username/password auth are not mixed.""" if (access_token or token_provider) and (username or password): raise ProgrammingError("Cannot use both token authentication and username/password") @@ -110,9 +116,9 @@ def create_client( password: str = "", access_token: str | None = None, token_provider: Callable[[], str] | None = None, - database: str = "__default__", + database: str | None = None, interface: str | None = None, - port: int = 0, + port: int | None = None, secure: bool | str = False, dsn: str | None = None, settings: dict[str, Any] | None = None, @@ -224,7 +230,7 @@ def create_client( interface, host, port, - username, + username or "", password, database, access_token, @@ -243,9 +249,9 @@ async def create_async_client( password: str = "", access_token: str | None = None, token_provider: Callable[[], str | Awaitable[str]] | None = None, - database: str = "__default__", + database: str | None = None, interface: str | None = None, - port: int = 0, + port: int | None = None, secure: bool | str = False, dsn: str | None = None, settings: dict[str, Any] | None = None, diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/driver/asyncclient.py b/contrib/python/clickhouse-connect/clickhouse_connect/driver/asyncclient.py index b0301df84ca..d7bc03884d4 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/driver/asyncclient.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/driver/asyncclient.py @@ -14,17 +14,18 @@ import uuid import zlib import zoneinfo from base64 import b64encode -from collections.abc import Awaitable, Callable, Generator, Iterable, Sequence +from collections.abc import Awaitable, Callable, Generator, Sequence from datetime import timezone, tzinfo from importlib import import_module from importlib.metadata import version as dist_version -from typing import TYPE_CHECKING, Any, BinaryIO +from typing import TYPE_CHECKING, Any, BinaryIO, cast import aiohttp import lz4.frame import zstandard if TYPE_CHECKING: + import numpy import pandas import polars import pyarrow @@ -64,6 +65,7 @@ from clickhouse_connect.driver.query import ( from clickhouse_connect.driver.streaming import StreamingFileAdapter, StreamingInsertSource, StreamingResponseSource from clickhouse_connect.driver.summary import QuerySummary from clickhouse_connect.driver.transform import NativeTransform +from clickhouse_connect.driver.types import Closable logger = logging.getLogger(__name__) columns_only_re = re.compile(r"LIMIT 0\s*$", re.IGNORECASE) @@ -319,7 +321,7 @@ class AsyncClient(Client): self._ssl_context = ssl_context self._proxy_url = proxy_url - self._connector_kwargs = { + self._connector_kwargs: dict[str, Any] = { "limit": connector_limit, "limit_per_host": connector_limit_per_host, "keepalive_timeout": keepalive_timeout, @@ -337,10 +339,10 @@ class AsyncClient(Client): self._read_format = "Native" self._write_format = "Native" self._transform = NativeTransform() - self._client_settings = {} + self._client_settings: dict[str, str] = {} self._initialized = False - self._reported_libs = set() - self._last_pool_reset = None + self._reported_libs: set[str] = set() + self._last_pool_reset: float | None = None self.headers["User-Agent"] = self.headers["User-Agent"].replace("mode:sync;", "mode:async;") if headers: self.headers.update(headers) @@ -505,18 +507,18 @@ class AsyncClient(Client): self._session = None raise - async def __aenter__(self): + async def __aenter__(self) -> AsyncClient: """Async context manager entry.""" if not self._initialized: await self._initialize() return self - async def __aexit__(self, exc_type, exc_val, exc_tb): + async def __aexit__(self, exc_type, exc_val, exc_tb) -> bool: """Async context manager exit.""" await self.close() return False - async def close(self): # type: ignore[override] + async def close(self) -> None: # type: ignore[override] async with self._session_lock: old_lease = self._session_lease self._session_lease = None @@ -524,7 +526,7 @@ class AsyncClient(Client): await old_lease.wait_drained() await old_lease.session.close() - async def close_connections(self): # type: ignore[override] + async def close_connections(self) -> None: # type: ignore[override] """Rotate the connection pool: new requests use a fresh session; in-flight requests keep using the old session until they complete, then it's closed.""" async with self._session_lock: @@ -543,7 +545,7 @@ class AsyncClient(Client): await old_lease.wait_drained() await old_lease.session.close() - def set_client_setting(self, key, value): + def set_client_setting(self, key: str, value: Any) -> None: str_value = self._validate_setting(key, value, common.get_setting("invalid_setting_action")) if str_value is not None: self._client_settings[key] = str_value @@ -555,12 +557,12 @@ class AsyncClient(Client): # Run sync providers off the event loop; await async providers. # The provider may be called concurrently if multiple requests get a 516 at the same time; # it must be safe to invoke in parallel (e.g. if it hits an IdP, consider rate limiting). - result = await asyncio.get_running_loop().run_in_executor(None, self._token_provider) + result = await asyncio.get_running_loop().run_in_executor(None, cast(Callable[[], str | Awaitable[str]], self._token_provider)) if inspect.isawaitable(result): result = await result return result - def set_access_token(self, access_token: str): + def set_access_token(self, access_token: str) -> None: auth_header = self.headers.get("Authorization") if auth_header and not auth_header.startswith("Bearer"): raise ProgrammingError("Cannot set access token when a different auth type is used") @@ -578,17 +580,18 @@ class AsyncClient(Client): return final_query + fmt async def _query_with_context(self, context: QueryContext) -> QueryResult: # type: ignore[override] - headers = {} - params = {} + headers: dict[str, Any] = {} + params: dict[str, str] = {} if self.database: params["database"] = self.database if self.protocol_version: - params["client_protocol_version"] = self.protocol_version + params["client_protocol_version"] = str(self.protocol_version) context.block_info = True params.update(self._validate_settings(context.settings)) context.rename_response_column = self._rename_response_column use_form = use_form_encoding(context.final_query, context.bind_params, self.form_encode_query_params) + files: dict[str, Any] | None = None if not context.is_insert and columns_only_re.search(context.uncommented_query): fmt_json_query = f"{context.final_query}\n FORMAT JSON" fields = {"query": fmt_json_query} @@ -641,7 +644,7 @@ class AsyncClient(Client): logger.debug("Failed to rename col '%s'. Skipping rename. Error: %s", name, e) names.append(name) types.append(get_from_name(col["type"])) - return QueryResult([], None, tuple(names), tuple(types)) + return QueryResult([], None, tuple(names), tuple(types)) # type: ignore[arg-type] if self.compression: headers["Accept-Encoding"] = self.compression @@ -651,7 +654,7 @@ class AsyncClient(Client): final_query = self._prep_query(context) files = None - data = None + data: Any = None if use_form: fields = {"query": final_query} @@ -755,7 +758,10 @@ class AsyncClient(Client): """ if query and query.lower().strip().startswith("select __connect_version__"): return QueryResult( - [[f"ClickHouse Connect v.{common.version()} â“’ ClickHouse Inc."]], None, ("connect_version",), (get_from_name("String"),) + [[f"ClickHouse Connect v.{common.version()} â“’ ClickHouse Inc."]], + None, # type: ignore[arg-type] # QueryContext.generator not yet Optional; widen after #805 merges + ("connect_version",), + (get_from_name("String"),), # type: ignore[arg-type] ) if not context: context = self.create_query_context( @@ -856,7 +862,7 @@ class AsyncClient(Client): """ return (await self._context_query(locals(), use_numpy=False, streaming=True)).rows_stream - async def query_np( + async def query_np( # type: ignore[override] self, query: str | None = None, parameters: Sequence | dict[str, Any] | None = None, @@ -869,7 +875,7 @@ class AsyncClient(Client): context: QueryContext | None = None, external_data: ExternalData | None = None, transport_settings: dict[str, str] | None = None, - ): + ) -> numpy.ndarray: check_numpy() self._add_integration_tag("numpy") return (await self._context_query(locals(), use_numpy=True)).np_result @@ -910,7 +916,7 @@ class AsyncClient(Client): use_extended_dtypes: bool | None = None, transport_settings: dict[str, str] | None = None, tz_mode: TzMode | None = None, - ): + ) -> pandas.DataFrame: check_pandas() self._add_integration_tag("pandas") return (await self._context_query(locals(), use_numpy=True, as_pandas=True)).df_result @@ -938,7 +944,7 @@ class AsyncClient(Client): self._add_integration_tag("pandas") return (await self._context_query(locals(), use_numpy=True, as_pandas=True, streaming=True)).df_stream - async def _context_query(self, lcls: dict, **overrides): # type: ignore[override] + async def _context_query(self, lcls: dict, **overrides): """ Helper method to create query context and execute query. Matches sync client pattern for consistency. @@ -983,6 +989,8 @@ class AsyncClient(Client): raise ProgrammingError("Command sent without query or recognized data") from None if payload or files: + if isinstance(cmd, bytes): + raise ProgrammingError("Binary parameter bind cannot be combined with command data or external data") from None params["query"] = cmd else: payload = cmd @@ -1104,7 +1112,7 @@ class AsyncClient(Client): async for chunk in response.content.iter_any(): yield chunk - class _RawStreamSource: + class _RawStreamSource(Closable): def close(self): try: response.close() @@ -1164,7 +1172,7 @@ class AsyncClient(Client): self, table: str | None = None, data: Sequence[Sequence[Any]] | None = None, - column_names: str | Iterable[str] = "*", + column_names: str | Sequence[str] | None = "*", database: str | None = None, column_types: Sequence[ClickHouseType] | None = None, column_type_names: Sequence[str] | None = None, @@ -1195,6 +1203,8 @@ class AsyncClient(Client): if (context is None or context.empty) and data is None: raise ProgrammingError("No data specified for insert") from None if context is None: + if table is None: + raise ProgrammingError("No table specified for insert") from None context = await self.create_insert_context( table, column_names, @@ -1219,7 +1229,7 @@ class AsyncClient(Client): use_strings: bool | None = None, external_data: ExternalData | None = None, transport_settings: dict[str, str] | None = None, - ): + ) -> pyarrow.Table: """ Query method using the ClickHouse Arrow format to return a PyArrow table :param query: Query statement/format string @@ -1312,9 +1322,9 @@ class AsyncClient(Client): streaming_source = StreamingResponseSource(response, encoding=encoding, exception_tag=exception_tag) await streaming_source.start_producer(loop) - queue = AsyncSyncQueue(maxsize=10) + queue: AsyncSyncQueue = AsyncSyncQueue(maxsize=10) - class _ArrowStreamSource: + class _ArrowStreamSource(Closable): def __init__(self, source, q): self._source = source self._queue = q @@ -1404,7 +1414,7 @@ class AsyncClient(Client): check_polars() self._add_integration_tag("polars") - def converter(table: pyarrow.Table) -> polars.DataFrame: + def converter(table: pyarrow.Table) -> polars.DataFrame: # type: ignore[misc] table = _apply_arrow_tz_policy(table, self.tz_mode) return options.pl.from_arrow(table) @@ -1458,7 +1468,7 @@ class AsyncClient(Client): check_polars() self._add_integration_tag("polars") - def converter(table: pyarrow.Table) -> polars.DataFrame: + def converter(table: pyarrow.Table) -> polars.DataFrame: # type: ignore[misc] table = _apply_arrow_tz_policy(table, self.tz_mode) return options.pl.from_arrow(table) @@ -1482,9 +1492,9 @@ class AsyncClient(Client): streaming_source = StreamingResponseSource(response, encoding=encoding, exception_tag=exception_tag) await streaming_source.start_producer(loop) - queue = AsyncSyncQueue(maxsize=10) + queue: AsyncSyncQueue = AsyncSyncQueue(maxsize=10) - class _ArrowDFStreamSource: + class _ArrowDFStreamSource(Closable): def __init__(self, source, q): self._source = source self._queue = q @@ -1560,7 +1570,7 @@ class AsyncClient(Client): column_names, insert_block = arrow_buffer(arrow_table, compression) if hasattr(insert_block, "to_pybytes"): insert_block = insert_block.to_pybytes() - return await self.raw_insert(full_table, column_names, insert_block, settings, "Arrow", transport_settings) + return await self.raw_insert(full_table, column_names, insert_block, settings, "Arrow", transport_settings=transport_settings) async def insert_df_arrow( # type: ignore[override] self, @@ -1719,8 +1729,8 @@ class AsyncClient(Client): active_source.start_producer() return active_source.async_generator() - headers = {"Content-Type": "application/octet-stream"} - if context.compression: + headers: dict[str, Any] = {"Content-Type": "application/octet-stream"} + if isinstance(context.compression, str): headers["Content-Encoding"] = context.compression params = {} @@ -1829,9 +1839,11 @@ class AsyncClient(Client): query = f"INSERT INTO {table}{cols} FORMAT {fmt_str}" if not compression and isinstance(insert_block, str): insert_block = query + "\n" + insert_block - elif not compression and isinstance(insert_block, (bytes, bytearray, BinaryIO)): + elif not compression and isinstance(insert_block, (bytes, bytearray)): insert_block = (query + "\n").encode() + insert_block else: + # Generators, file-like objects, and compressed data: send the + # INSERT query as a URL param and stream the body as-is. params["query"] = query if self.database: @@ -2040,7 +2052,7 @@ class AsyncClient(Client): response = await session.request(**request_kwargs) if 200 <= response.status < 300 and not response.headers.get(ex_header): # Caller releases lease after consuming the body. - response._lease_release = _one_shot(lease.release) + response._lease_release = _one_shot(lease.release) # type: ignore[attr-defined] lease_released = True return response diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/driver/asyncqueue.py b/contrib/python/clickhouse-connect/clickhouse_connect/driver/asyncqueue.py index a318f59cea1..3d1c277ebd6 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/driver/asyncqueue.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/driver/asyncqueue.py @@ -1,7 +1,7 @@ import asyncio import threading from collections import deque -from typing import Generic, TypeVar +from typing import Any, Generic, TypeVar from clickhouse_connect.driver.exceptions import ProgrammingError @@ -9,7 +9,8 @@ __all__ = ["AsyncSyncQueue", "Empty", "Full", "EOF_SENTINEL"] T = TypeVar("T") -EOF_SENTINEL = object() +# Typed Any so it can stand in for a queued T value on EOF without a cast in every reader. +EOF_SENTINEL: Any = object() class AsyncSyncQueue(Generic[T]): @@ -73,7 +74,8 @@ class AsyncSyncQueue(Generic[T]): while waiter_queue: fut = waiter_queue.popleft() if not fut.done(): - self._loop.call_soon_threadsafe(self._safe_set_result, fut) + # _bind_loop() runs before any Future is created, so _loop is always set here + self._loop.call_soon_threadsafe(self._safe_set_result, fut) # type: ignore[union-attr] break def shutdown(self): @@ -161,7 +163,8 @@ class _AsyncQueueInterface(Generic[T]): if self._p._shutdown: return EOF_SENTINEL - fut = self._p._loop.create_future() + # _bind_loop() is called at the top of get(), so _loop is always set here + fut = self._p._loop.create_future() # type: ignore[union-attr] self._p._async_getters.append(fut) try: @@ -185,7 +188,8 @@ class _AsyncQueueInterface(Generic[T]): self._p._wakeup_async_waiter(self._p._async_getters) return - fut = self._p._loop.create_future() + # _bind_loop() is called at the top of put(), so _loop is always set here + fut = self._p._loop.create_future() # type: ignore[union-attr] self._p._async_putters.append(fut) try: diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/driver/binding.py b/contrib/python/clickhouse-connect/clickhouse_connect/driver/binding.py index b4f848dc77a..52fe6333ce7 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/driver/binding.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/driver/binding.py @@ -23,7 +23,7 @@ class DT64Param: def __init__(self, value: datetime): self.value = value - def format(self, tz: tzinfo, top_level: bool) -> str: + def format(self, tz: tzinfo | None, top_level: bool) -> str: value = self.value if tz: value = value.astimezone(tz) @@ -128,12 +128,13 @@ def bind_query( query: str, parameters: Sequence | dict[str, Any] | None, server_tz: tzinfo | None = None, -) -> tuple[str, dict[str, str]]: +) -> tuple[str | bytes, dict[str, str]]: query = query.rstrip(";") if not parameters: return query, {} binary_binds = None + bound_params: dict[str, str] = {} if isinstance(parameters, dict): params_copy = dict_copy(parameters) @@ -186,13 +187,14 @@ def bind_query( break binary_indexes[item_index + len(key)] = key, v item_index += len(key) - query = b"" + binary_out = b"" start = 0 for loc in sorted(binary_indexes.keys()): key, value = binary_indexes[loc] - query += binary_query[start:loc] + value + key + binary_out += binary_query[start:loc] + value + key start = loc - query += binary_query[start:] + binary_out += binary_query[start:] + return binary_out, bound_params return query, bound_params @@ -232,7 +234,7 @@ def escape_bytes(value): return "".join(f"{BS}x{b:02x}" for b in value) -def format_query_value(value: Any, server_tz: tzinfo = timezone.utc): +def format_query_value(value: Any, server_tz: tzinfo | None = timezone.utc): """ Format Python values in a ClickHouse query :param value: Python object @@ -269,11 +271,11 @@ def format_query_value(value: Any, server_tz: tzinfo = timezone.utc): return value -def str_query_value(value: Any, server_tz: tzinfo = timezone.utc): +def str_query_value(value: Any, server_tz: tzinfo | None = timezone.utc): return str(format_query_value(value, server_tz)) -def format_bind_value(value: Any, server_tz: tzinfo = timezone.utc, top_level: bool = True): +def format_bind_value(value: Any, server_tz: tzinfo | None = timezone.utc, top_level: bool = True): """ Format Python values in a ClickHouse query :param value: Python object diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/driver/buffer.py b/contrib/python/clickhouse-connect/clickhouse_connect/driver/buffer.py index 97c01670b09..c0da348660f 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/driver/buffer.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/driver/buffer.py @@ -114,11 +114,11 @@ class ResponseBuffer(ByteSource): def read_str_col( self, num_rows: int, - encoding: str, + encoding: str | None, nullable: bool = False, null_obj: Any = None, ) -> Iterable[str]: - column = [] + column: list[Any] = [] app = column.append null_map = self.read_bytes(num_rows) if nullable else None for ix in range(num_rows): @@ -148,7 +148,7 @@ class ResponseBuffer(ByteSource): def read_fixed_str_col(self, sz: int, num_rows: int, encoding: str) -> Iterable[str]: source = self.read_bytes(sz * num_rows) - column = [] + column: list[str] = [] app = column.append for ix in range(0, sz * num_rows, sz): try: @@ -167,7 +167,7 @@ class ResponseBuffer(ByteSource): return column @property - def last_message(self) -> bytes: + def last_message(self) -> bytes | None: # type: ignore[override] # overrides writable attr with property return self.buffer def close(self): diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/driver/client.py b/contrib/python/clickhouse-connect/clickhouse_connect/driver/client.py index e5da4ec5736..46ce2b4b76a 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/driver/client.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/driver/client.py @@ -3,12 +3,13 @@ from __future__ import annotations import io import logging from abc import ABC, abstractmethod -from collections.abc import Generator, Iterable, Sequence +from collections.abc import Generator, Sequence from datetime import timezone, tzinfo from typing import ( TYPE_CHECKING, Any, BinaryIO, + cast, ) from zoneinfo import ZoneInfoNotFoundError @@ -55,6 +56,7 @@ from clickhouse_connect.driver.query import ( to_arrow_batches, ) from clickhouse_connect.driver.summary import QuerySummary +from clickhouse_connect.driver.types import Closable if TYPE_CHECKING: import numpy @@ -62,7 +64,7 @@ if TYPE_CHECKING: import polars import pyarrow -io.DEFAULT_BUFFER_SIZE = 1024 * 256 +io.DEFAULT_BUFFER_SIZE = 1024 * 256 # type: ignore[misc] # override module default buffer size logger = logging.getLogger(__name__) arrow_str_setting = "output_format_arrow_string_as_string" @@ -117,8 +119,8 @@ class Client(ABC): compression: str | None = None write_compression: str | None = None protocol_version = 0 - valid_transport_settings = set() - optional_transport_settings = set() + valid_transport_settings: set[str] = set() + optional_transport_settings: set[str] = set() database = None max_error_message = 0 _tz_source: TzSource = "auto" @@ -169,7 +171,7 @@ class Client(ABC): """ self.query_limit = coerce_int(query_limit) self.query_retries = coerce_int(query_retries) - if database and not database == "__default__": + if database and database != "__default__": self.database = database if show_clickhouse_errors is not None: self.show_clickhouse_errors = coerce_bool(show_clickhouse_errors) @@ -184,9 +186,9 @@ class Client(ABC): self._tz_source = resolved_tz_source # Initialize attributes that will be set during connection - self.server_version = None - self.server_tz = timezone.utc - self.server_settings = {} + self.server_version: str | None = None + self.server_tz: tzinfo = timezone.utc + self.server_settings: dict[str, SettingDef] = {} if autoconnect: self._init_common_settings(resolved_tz_source) @@ -196,7 +198,10 @@ class Client(ABC): def _init_common_settings(self, tz_source: TzSource): self.server_tz, self._dst_safe = timezone.utc, True - self.server_version, server_tz = tuple(self.command("SELECT version(), timezone()", use_database=False)) + version_result = self.command("SELECT version(), timezone()", use_database=False) + if not isinstance(version_result, Sequence) or isinstance(version_result, str): + raise OperationalError(f"Unexpected response to server version query: {version_result!r}") + self.server_version, server_tz = version_result[0], version_result[1] try: server_tz_info = tzutil.resolve_zone(server_tz) server_tz_info, self._dst_safe = tzutil.normalize_timezone(server_tz_info, trust_fixed_offset=True) @@ -247,8 +252,10 @@ class Client(ABC): :param settings: Dictionary of setting name and values :return: A filtered dictionary of settings with values rendered as strings """ - validated = {} + validated: dict[str, str] = {} invalid_action = common.get_setting("invalid_setting_action") + if not settings: + return validated for key, value in settings.items(): str_value = self._validate_setting(key, value, invalid_action) if str_value is not None: @@ -349,7 +356,7 @@ class Client(ABC): column_oriented: bool | None = None, use_numpy: bool | None = None, max_str_len: int | None = None, - context: QueryContext = None, + context: QueryContext | None = None, query_tz: str | tzinfo | None = None, column_tzs: dict[str, str | tzinfo] | None = None, external_data: ExternalData | None = None, @@ -363,14 +370,17 @@ class Client(ABC): """ if query and query.lower().strip().startswith("select __connect_version__"): return QueryResult( - [[f"ClickHouse Connect v.{version()} â“’ ClickHouse Inc."]], None, ("connect_version",), (get_from_name("String"),) + [[f"ClickHouse Connect v.{version()} â“’ ClickHouse Inc."]], + None, + ("connect_version",), + (get_from_name("String"),), ) kwargs = locals().copy() del kwargs["self"] query_context = self.create_query_context(**kwargs) if query_context.is_command: response = self.command( - query, + cast(str, query), parameters=query_context.parameters, settings=query_context.settings, external_data=query_context.external_data, @@ -390,7 +400,7 @@ class Client(ABC): column_formats: dict[str, str | dict[str, str]] | None = None, encoding: str | None = None, use_none: bool | None = None, - context: QueryContext = None, + context: QueryContext | None = None, query_tz: str | tzinfo | None = None, column_tzs: dict[str, str | tzinfo] | None = None, external_data: ExternalData | None = None, @@ -413,7 +423,7 @@ class Client(ABC): column_formats: dict[str, str | dict[str, str]] | None = None, encoding: str | None = None, use_none: bool | None = None, - context: QueryContext = None, + context: QueryContext | None = None, query_tz: str | tzinfo | None = None, column_tzs: dict[str, str | tzinfo] | None = None, external_data: ExternalData | None = None, @@ -436,7 +446,7 @@ class Client(ABC): column_formats: dict[str, str | dict[str, str]] | None = None, encoding: str | None = None, use_none: bool | None = None, - context: QueryContext = None, + context: QueryContext | None = None, query_tz: str | tzinfo | None = None, column_tzs: dict[str, str | tzinfo] | None = None, external_data: ExternalData | None = None, @@ -456,7 +466,7 @@ class Client(ABC): query: str, parameters: Sequence | dict[str, Any] | None = None, settings: dict[str, Any] | None = None, - fmt: str = None, + fmt: str | None = None, use_database: bool = True, external_data: ExternalData | None = None, transport_settings: dict[str, str] | None = None, @@ -480,7 +490,7 @@ class Client(ABC): query: str, parameters: Sequence | dict[str, Any] | None = None, settings: dict[str, Any] | None = None, - fmt: str = None, + fmt: str | None = None, use_database: bool = True, external_data: ExternalData | None = None, transport_settings: dict[str, str] | None = None, @@ -508,7 +518,7 @@ class Client(ABC): encoding: str | None = None, use_none: bool | None = None, max_str_len: int | None = None, - context: QueryContext = None, + context: QueryContext | None = None, external_data: ExternalData | None = None, transport_settings: dict[str, str] | None = None, ) -> numpy.ndarray: @@ -531,7 +541,7 @@ class Client(ABC): encoding: str | None = None, use_none: bool | None = None, max_str_len: int | None = None, - context: QueryContext = None, + context: QueryContext | None = None, external_data: ExternalData | None = None, transport_settings: dict[str, str] | None = None, ) -> StreamContext: @@ -557,7 +567,7 @@ class Client(ABC): use_na_values: bool | None = None, query_tz: str | None = None, column_tzs: dict[str, str | tzinfo] | None = None, - context: QueryContext = None, + context: QueryContext | None = None, external_data: ExternalData | None = None, use_extended_dtypes: bool | None = None, transport_settings: dict[str, str] | None = None, @@ -585,7 +595,7 @@ class Client(ABC): use_na_values: bool | None = None, query_tz: str | None = None, column_tzs: dict[str, str | tzinfo] | None = None, - context: QueryContext = None, + context: QueryContext | None = None, external_data: ExternalData | None = None, use_extended_dtypes: bool | None = None, transport_settings: dict[str, str] | None = None, @@ -688,7 +698,7 @@ class Client(ABC): if as_pandas and use_extended_dtypes is None: use_extended_dtypes = True return QueryContext( - query=query, + query=cast(str | bytes, query), parameters=parameters, settings=settings, query_formats=query_formats, @@ -766,13 +776,16 @@ class Client(ABC): self._add_integration_tag("arrow") settings = self._update_arrow_settings(settings, use_strings) return to_arrow_batches( - self.raw_stream( - query, - parameters, - settings, - fmt="ArrowStream", - external_data=external_data, - transport_settings=transport_settings, + cast( + io.IOBase, + self.raw_stream( + query, + parameters, + settings, + fmt="ArrowStream", + external_data=external_data, + transport_settings=transport_settings, + ), ) ) @@ -814,7 +827,7 @@ class Client(ABC): check_polars() self._add_integration_tag("polars") - def converter(table: pyarrow.Table) -> polars.DataFrame: + def converter(table: pyarrow.Table) -> polars.DataFrame: # type: ignore[misc] table = _apply_arrow_tz_policy(table, self.tz_mode) return options.pl.from_arrow(table) @@ -867,7 +880,7 @@ class Client(ABC): check_polars() self._add_integration_tag("polars") - def converter(table: pyarrow.Table) -> polars.DataFrame: + def converter(table: pyarrow.Table) -> polars.DataFrame: # type: ignore[misc] table = _apply_arrow_tz_policy(table, self.tz_mode) return options.pl.from_arrow(table) else: @@ -882,7 +895,7 @@ class Client(ABC): for batch in reader: yield converter(batch) - return StreamContext(raw_stream, df_generator()) + return StreamContext(cast(Closable, raw_stream), df_generator()) def _update_arrow_settings(self, settings: dict[str, Any] | None, use_strings: bool | None) -> dict[str, Any]: settings = dict_copy(settings) @@ -903,8 +916,8 @@ class Client(ABC): self, cmd: str, parameters: Sequence | dict[str, Any] | None = None, - data: str | bytes = None, - settings: dict[str, Any] = None, + data: str | bytes | None = None, + settings: dict[str, Any] | None = None, use_database: bool = True, external_data: ExternalData | None = None, transport_settings: dict[str, str] | None = None, @@ -934,14 +947,14 @@ class Client(ABC): def insert( self, table: str | None = None, - data: Sequence[Sequence[Any]] = None, - column_names: str | Iterable[str] = "*", + data: Sequence[Sequence[Any]] | None = None, + column_names: str | Sequence[str] | None = "*", database: str | None = None, - column_types: Sequence[ClickHouseType] = None, - column_type_names: Sequence[str] = None, + column_types: Sequence[ClickHouseType] | None = None, + column_type_names: Sequence[str] | None = None, column_oriented: bool = False, settings: dict[str, Any] | None = None, - context: InsertContext = None, + context: InsertContext | None = None, transport_settings: dict[str, str] | None = None, ) -> QuerySummary: """ @@ -966,6 +979,8 @@ class Client(ABC): if (context is None or context.empty) and data is None: raise ProgrammingError("No data specified for insert") from None if context is None: + if table is None: + raise ProgrammingError("No table specified for insert") from None context = self.create_insert_context( table, column_names, @@ -984,14 +999,14 @@ class Client(ABC): def insert_df( self, - table: str = None, + table: str | None = None, df=None, database: str | None = None, settings: dict | None = None, column_names: Sequence[str] | None = None, - column_types: Sequence[ClickHouseType] = None, - column_type_names: Sequence[str] = None, - context: InsertContext = None, + column_types: Sequence[ClickHouseType] | None = None, + column_type_names: Sequence[str] | None = None, + context: InsertContext | None = None, transport_settings: dict[str, str] | None = None, ) -> QuerySummary: """ @@ -1034,7 +1049,7 @@ class Client(ABC): self, table: str, arrow_table, - database: str = None, + database: str | None = None, settings: dict | None = None, transport_settings: dict[str, str] | None = None, ) -> QuerySummary: @@ -1051,7 +1066,7 @@ class Client(ABC): full_table = table if "." in table or not database else f"{database}.{table}" compression = self.write_compression if self.write_compression in ("zstd", "lz4") else None column_names, insert_block = arrow_buffer(arrow_table, compression) - return self.raw_insert(full_table, column_names, insert_block, settings, "Arrow", transport_settings) + return self.raw_insert(full_table, column_names, insert_block, settings, "Arrow", transport_settings=transport_settings) def insert_df_arrow( self, @@ -1117,8 +1132,8 @@ class Client(ABC): table: str, column_names: str | Sequence[str] | None = None, database: str | None = None, - column_types: Sequence[ClickHouseType] = None, - column_type_names: Sequence[str] = None, + column_types: Sequence[ClickHouseType] | None = None, + column_type_names: Sequence[str] | None = None, column_oriented: bool = False, settings: dict[str, Any] | None = None, data: Sequence[Sequence[Any]] | None = None, @@ -1146,7 +1161,7 @@ class Client(ABC): full_table = f"{quote_identifier(database)}.{quote_identifier(table)}" else: full_table = quote_identifier(table) - column_defs = [] + column_defs: list[ColumnDef] = [] if column_types is None and column_type_names is None: describe_result = self.query(f"DESCRIBE TABLE {full_table}", settings=settings) column_defs = [ @@ -1189,7 +1204,7 @@ class Client(ABC): :return: True if version_str is greater than the server_version, False if less than """ try: - server_parts = [int(x) for x in self.server_version.split(".") if x.isnumeric()] + server_parts = [int(x) for x in (self.server_version or "").split(".") if x.isnumeric()] server_parts.extend([0] * (4 - len(server_parts))) version_parts = [int(x) for x in version_str.split(".")] version_parts.extend([0] * (4 - len(version_parts))) @@ -1222,7 +1237,7 @@ class Client(ABC): self, table: str, column_names: Sequence[str] | None = None, - insert_block: str | bytes | Generator[bytes, None, None] | BinaryIO = None, + insert_block: str | bytes | Generator[bytes, None, None] | BinaryIO | None = None, settings: dict | None = None, fmt: str | None = None, compression: str | None = None, diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/driver/common.py b/contrib/python/clickhouse-connect/clickhouse_connect/driver/common.py index a4acbd292df..fc609bfd84b 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/driver/common.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/driver/common.py @@ -3,6 +3,7 @@ import asyncio import struct import sys from collections.abc import Callable, Generator, MutableSequence, Sequence +from io import IOBase from typing import Any from clickhouse_connect.driver.exceptions import DataError, ProgrammingError, StreamClosedError @@ -107,7 +108,7 @@ def unescape_identifier(x: str) -> str: return x -def dict_copy(source: dict = None, update: dict | None = None) -> dict: +def dict_copy(source: dict | None = None, update: dict | None = None) -> dict: copy = source.copy() if source else {} if update: copy.update(update) @@ -153,10 +154,13 @@ class SliceView(Sequence): slots = ("_source", "_range") + _source: Sequence + _range: range + def __init__(self, source: Sequence, source_slice: slice | None = None): if isinstance(source, SliceView): self._source = source._source - self._range = source._range[source_slice] + self._range = source._range if source_slice is None else source._range[source_slice] else: self._source = source if source_slice is None: @@ -200,7 +204,7 @@ class StreamContext: __slots__ = "source", "gen", "_in_context" - def __init__(self, source: Closable, gen: Generator): + def __init__(self, source: Closable | IOBase, gen: Generator): self.source = source self.gen = gen self._in_context = False diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/driver/compression.py b/contrib/python/clickhouse-connect/clickhouse_connect/driver/compression.py index 802b6e815c2..b6533a95711 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/driver/compression.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/driver/compression.py @@ -1,5 +1,4 @@ import zlib -from abc import abstractmethod import lz4 import lz4.frame @@ -17,14 +16,13 @@ if brotli: available_compression.append("br") available_compression.extend(["gzip", "deflate"]) -comp_map = {} +comp_map: dict[str, "Compressor | type[Compressor]"] = {} class Compressor: def __init_subclass__(cls, tag: str, thread_safe: bool = True): comp_map[tag] = cls() if thread_safe else cls - @abstractmethod def compress_block(self, block) -> bytes | bytearray: return block @@ -66,11 +64,10 @@ class BrotliCompressor(Compressor, tag="br"): null_compressor = Compressor() -def get_compressor(compression: str) -> Compressor: +def get_compressor(compression: str | None) -> Compressor: if not compression: return null_compressor comp = comp_map[compression] - try: - return comp() - except TypeError: + if isinstance(comp, Compressor): return comp + return comp() diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/driver/context.py b/contrib/python/clickhouse-connect/clickhouse_connect/driver/context.py index 50b8315a0db..eec0ed517ed 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/driver/context.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/driver/context.py @@ -5,7 +5,7 @@ from typing import Any logger = logging.getLogger(__name__) -_empty_map = {} +_empty_map: dict[Any, Any] = {} class BaseQueryContext: @@ -38,7 +38,7 @@ class BaseQueryContext: self.query_formats = query_formats or {} self.column_formats = column_formats or {} self.transport_settings = transport_settings - self.column_name = None + self.column_name: str | None = None self.encoding = encoding self.use_numpy = use_numpy self.use_extended_dtypes = use_extended_dtypes diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/driver/dataconv.py b/contrib/python/clickhouse-connect/clickhouse_connect/driver/dataconv.py index 5a6d1786f8a..9de76e3fbdd 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/driver/dataconv.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/driver/dataconv.py @@ -17,11 +17,14 @@ MONTH_DAYS_LEAP = (0, 31, 60, 91, 121, 152, 182, 213, 244, 274, 305, 335, 366) def read_ipv4_col(source: ByteSource, num_rows: int): column = source.read_array("I", num_rows) fast_ip_v4 = IPv4Address.__new__ - new_col = [] + new_col: list[IPv4Address] = [] app = new_col.append for x in column: ipv4 = fast_ip_v4(IPv4Address) - ipv4._ip = x + # _ip is CPython's private backing int for the address. + # It's directly set to bypass IPv4Address.__init__ + # for speed when bulk-decoding a column + ipv4._ip = x # type: ignore[attr-defined] app(ipv4) return new_col @@ -103,7 +106,7 @@ def read_uuid_col(source: ByteSource, num_rows: int): new_uuid = UUID.__new__ unsafe = SafeUUID.unsafe oset = object.__setattr__ - column = [] + column: list[UUID] = [] app = column.append for i in range(num_rows): ix = i << 1 @@ -173,7 +176,7 @@ def write_str_col(column: Sequence, nullable: bool, encoding: str | None, dest: return 0 -def write_native_col(code: str, column: Sequence, dest: bytearray, col_name: object = None) -> int: +def write_native_col(code: str, column: Sequence, dest: bytearray, col_name: str | None = None) -> int: """ Pure Python fallback for write_native_col. Delegates to write_array which uses struct.pack. diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/driver/ddl.py b/contrib/python/clickhouse-connect/clickhouse_connect/driver/ddl.py index d6137d124b4..b682fbfa4e6 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/driver/ddl.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/driver/ddl.py @@ -1,24 +1,30 @@ from __future__ import annotations from collections.abc import Sequence -from typing import TYPE_CHECKING, NamedTuple +from typing import TYPE_CHECKING, NamedTuple, Protocol -from clickhouse_connect.datatypes.base import ClickHouseType from clickhouse_connect.driver.options import check_arrow if TYPE_CHECKING: import pyarrow +class _NamedType(Protocol): + """Structural type for anything exposing a ClickHouse type ``name``.""" + + @property + def name(self) -> str: ... + + class TableColumnDef(NamedTuple): """ Simplified ClickHouse Table Column definition for DDL """ name: str - ch_type: ClickHouseType - expr_type: str = None - expr: str = None + ch_type: _NamedType + expr_type: str | None = None + expr: str | None = None @property def col_expr(self): diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/driver/exceptions.py b/contrib/python/clickhouse-connect/clickhouse_connect/driver/exceptions.py index cddcac55f59..f00ff9aa712 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/driver/exceptions.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/driver/exceptions.py @@ -29,7 +29,7 @@ class ClickHouseError(Exception): """Exception related to operation with ClickHouse.""" -class Warning(Warning, ClickHouseError): # noqa: N818 +class Warning(Warning, ClickHouseError): # type: ignore[misc] # noqa: N818 """Exception raised for important warnings like data truncations while inserting, etc.""" diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/driver/httpclient.py b/contrib/python/clickhouse-connect/clickhouse_connect/driver/httpclient.py index 170c2ab6ab8..9d1aa98b0da 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/driver/httpclient.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/driver/httpclient.py @@ -8,7 +8,7 @@ from base64 import b64encode from collections.abc import Callable, Generator, Sequence from importlib import import_module from importlib.metadata import version as dist_version -from typing import Any, BinaryIO +from typing import Any, BinaryIO, cast from urllib.parse import urlencode from urllib3 import Timeout @@ -43,7 +43,7 @@ from clickhouse_connect.driver.httputil import ( get_response_data, ) from clickhouse_connect.driver.insert import InsertContext -from clickhouse_connect.driver.query import QueryContext, QueryResult, TzSource, returns_empty_string_on_empty_body +from clickhouse_connect.driver.query import QueryContext, QueryResult, TzMode, TzSource, returns_empty_string_on_empty_body from clickhouse_connect.driver.summary import QuerySummary from clickhouse_connect.driver.transform import NativeTransform @@ -57,7 +57,7 @@ _REMOTE_CLOSE_ERRORS = (ConnectionResetError, BrokenPipeError) class HttpClient(Client): - params = {} + params: dict[str, str] = {} valid_transport_settings = { "database", "buffer_size", @@ -84,7 +84,7 @@ class HttpClient(Client): port: int, username: str, password: str, - database: str, + database: str | None, access_token: str | None = None, token_provider: Callable[[], str] | None = None, compress: bool | str = True, @@ -141,7 +141,7 @@ class HttpClient(Client): self.headers["X-ClickHouse-SSL-Certificate-Auth"] = "on" if not self.http and (server_host_name or ca_cert or client_cert or not verify or https_proxy): - options = {"verify": verify} + options: dict[str, Any] = {"verify": verify} dict_add(options, "ca_cert", ca_cert) dict_add(options, "client_cert", client_cert) dict_add(options, "client_cert_key", client_cert_key) @@ -167,7 +167,7 @@ class HttpClient(Client): elif (not client_cert or tls_mode in ("strict", "proxy")) and username: self.headers["Authorization"] = "Basic " + b64encode(f"{username}:{password}".encode()).decode() - self._reported_libs = set() + self._reported_libs: set[str] = set() self.headers["User-Agent"] = common.build_client_name(client_name) if headers: self.headers.update(headers) @@ -184,7 +184,7 @@ class HttpClient(Client): self._send_progress = None self._send_comp_setting = False self._progress_interval = None - self._active_session = None + self._active_session: str | None = None self._rename_response_column = rename_response_column # allow to override the global autogenerate_session_id setting via the constructor params @@ -220,7 +220,7 @@ class HttpClient(Client): query_retries=query_retries, server_host_name=server_host_name, tz_source=tz_source, - tz_mode=tz_mode, + tz_mode=cast(TzMode | None, tz_mode), show_clickhouse_errors=show_clickhouse_errors, autoconnect=True, ) @@ -265,16 +265,17 @@ class HttpClient(Client): return final_query + fmt def _query_with_context(self, context: QueryContext) -> QueryResult: - headers = {} - params = {} + headers: dict[str, Any] = {} + params: dict[str, str] = {} if self.database: params["database"] = self.database if self.protocol_version: - params["client_protocol_version"] = self.protocol_version + params["client_protocol_version"] = str(self.protocol_version) context.block_info = True params.update(self._validate_settings(context.settings)) context.rename_response_column = self._rename_response_column use_form = use_form_encoding(context.final_query, context.bind_params, self.form_encode_query_params) + fields: dict[str, Any] | None = None if not context.is_insert and columns_only_re.search(context.uncommented_query): # Mirror normal query behavior for form encoding and external data fmt_json_query = f"{context.final_query}\n FORMAT JSON" @@ -309,7 +310,7 @@ class HttpClient(Client): logger.debug("Failed to rename col '%s'. Skipping rename. Error: %s", name, e) names.append(name) types.append(registry.get_from_name(col["type"])) - return QueryResult([], None, tuple(names), tuple(types)) + return QueryResult([], None, tuple(names), tuple(types)) # type: ignore[arg-type] if self.compression: headers["Accept-Encoding"] = self.compression @@ -347,10 +348,12 @@ class HttpClient(Client): ) exception_tag = response.headers.get(ex_tag_header) byte_source = RespBuffCls(ResponseSource(response, exception_tag=exception_tag)) - context.set_response_tz(self._check_tz_change(response.headers.get("X-ClickHouse-Timezone"))) + response_tz = self._check_tz_change(response.headers.get("X-ClickHouse-Timezone")) + if response_tz is not None: + context.set_response_tz(response_tz) query_result = self._transform.parse_response(byte_source, context) query_result.summary = self._summary(response) - return query_result + return cast(QueryResult, query_result) def data_insert(self, context: InsertContext) -> QuerySummary: """ @@ -371,7 +374,7 @@ class HttpClient(Client): headers = {"Content-Type": "application/octet-stream"} if context.compression is None: context.compression = self.write_compression - if context.compression: + if isinstance(context.compression, str): headers["Content-Encoding"] = context.compression block_gen = self._transform.build_insert(context) @@ -401,9 +404,9 @@ class HttpClient(Client): def raw_insert( self, - table: str = None, + table: str | None = None, column_names: Sequence[str] | None = None, - insert_block: str | bytes | Generator[bytes, None, None] | BinaryIO = None, + insert_block: str | bytes | Generator[bytes, None, None] | BinaryIO | None = None, settings: dict | None = None, fmt: str | None = None, compression: str | None = None, @@ -421,9 +424,11 @@ class HttpClient(Client): query = f"INSERT INTO {table}{cols} FORMAT {fmt if fmt else self._write_format}" if not compression and isinstance(insert_block, str): insert_block = query + "\n" + insert_block - elif not compression and isinstance(insert_block, (bytes, bytearray, BinaryIO)): + elif not compression and isinstance(insert_block, (bytes, bytearray)): insert_block = (query + "\n").encode() + insert_block else: + # Generators, file-like objects, and compressed data: send the + # INSERT query as a URL param and stream the body as-is. params["query"] = query if self.database: params["database"] = self.database @@ -448,7 +453,7 @@ class HttpClient(Client): self, cmd: str, parameters: Sequence | dict[str, Any] | None = None, - data: str | bytes = None, + data: str | bytes | None = None, settings: dict | None = None, use_database: bool = True, external_data: ExternalData | None = None, @@ -457,9 +462,9 @@ class HttpClient(Client): """ See BaseClient doc_string for this method """ - cmd, params = bind_query(cmd, parameters, self.server_tz) - headers = {} - payload = None + bound_cmd, params = bind_query(cmd, parameters, self.server_tz) + headers: dict[str, Any] = {} + payload: str | bytes | None = None fields = None if external_data: if data: @@ -472,12 +477,14 @@ class HttpClient(Client): elif isinstance(data, bytes): headers["Content-Type"] = "application/octet-stream" payload = data - if payload is None and not cmd: + if payload is None and not bound_cmd: raise ProgrammingError("Command sent without query or recognized data") from None if payload or fields: - params["query"] = cmd + if isinstance(bound_cmd, bytes): + raise ProgrammingError("Binary parameter bind cannot be combined with command data or external data") from None + params["query"] = bound_cmd else: - payload = cmd + payload = bound_cmd if use_database and self.database: params["database"] = self.database params.update(self._validate_settings(settings or {})) @@ -495,7 +502,7 @@ class HttpClient(Client): return result except UnicodeDecodeError: return str(response.data) - if returns_empty_string_on_empty_body(cmd): + if returns_empty_string_on_empty_body(bound_cmd): return "" return QuerySummary(self._summary(response)) @@ -546,7 +553,7 @@ class HttpClient(Client): stream: bool = False, server_wait: bool = True, fields: dict[str, tuple] | None = None, - error_handler: Callable = None, + error_handler: Callable | None = None, retry_body: Callable[[], Any] | None = None, ) -> HTTPResponse: if isinstance(data, str): @@ -571,7 +578,7 @@ class HttpClient(Client): final_params["query_id"] = str(uuid.uuid4()) url = f"{self.url}?{urlencode(final_params)}" - kwargs = {"headers": headers, "timeout": self.timeout, "retries": self.http_retries, "preload_content": not stream} + kwargs: dict[str, Any] = {"headers": headers, "timeout": self.timeout, "retries": self.http_retries, "preload_content": not stream} if self.server_host_name: kwargs["assert_same_host"] = False kwargs["headers"].update({"Host": self.server_host_name}) @@ -579,7 +586,7 @@ class HttpClient(Client): kwargs["fields"] = fields else: kwargs["body"] = data - check_conn_expiration(self.http) + check_conn_expiration(cast(PoolManager, self.http)) query_session = final_params.get("session_id") while True: attempts += 1 @@ -593,7 +600,7 @@ class HttpClient(Client): # throw an error instead, but in most cases this more helpful error will be thrown first self._active_session = query_session try: - response = self.http.request(method, url, **kwargs) + response: HTTPResponse = cast(HTTPResponse, cast(PoolManager, self.http).request(method, url, **kwargs)) except HTTPError as ex: # Always allow at least one retry on a clean connection error so a single stale # keep-alive socket doesn't surface to the caller, and additionally honor the @@ -639,7 +646,7 @@ class HttpClient(Client): kwargs["body"] = retry_body() response.close() logger.debug("Refreshing access token after authentication failure") - elif error_handler: + elif error_handler is not None: error_handler(response) else: self._error_handler(response) @@ -649,7 +656,7 @@ class HttpClient(Client): query: str, parameters: Sequence | dict[str, Any] | None = None, settings: dict[str, Any] | None = None, - fmt: str = None, + fmt: str | None = None, use_database: bool = True, external_data: ExternalData | None = None, transport_settings: dict[str, str] | None = None, @@ -665,7 +672,7 @@ class HttpClient(Client): query: str, parameters: Sequence | dict[str, Any] | None = None, settings: dict[str, Any] | None = None, - fmt: str = None, + fmt: str | None = None, use_database: bool = True, external_data: ExternalData | None = None, transport_settings: dict[str, str] | None = None, @@ -689,7 +696,7 @@ class HttpClient(Client): query: str, parameters: Sequence | dict[str, Any] | None, settings: dict[str, Any] | None, - fmt: str, + fmt: str | None, use_database: bool, external_data: ExternalData | None, ): @@ -699,22 +706,24 @@ class HttpClient(Client): params = self._validate_settings(settings or {}) if use_database and self.database: params["database"] = self.database - fields = {} + form_fields: dict[str, Any] = {} + fields: dict[str, Any] | None = form_fields use_form = use_form_encoding(final_query, bind_params, self.form_encode_query_params) # Setup query body if external_data and not use_form and isinstance(final_query, bytes): raise ProgrammingError("Binary query cannot be placed in URL when using External Data; enable form encoding.") # Setup additional query parameters and body + body: str | bytes = b"" if use_form: - body = b"" - fields["query"] = final_query - fields.update(bind_params) + form_fields["query"] = final_query + form_fields.update(bind_params) if external_data: params.update(external_data.query_params) - fields.update(external_data.form_data) + form_fields.update(external_data.form_data) elif external_data: params.update(bind_params) - body = b"" + # Guaranteed str: the check above raises if external_data and not use_form and bytes + assert isinstance(final_query, str) params["query"] = final_query params.update(external_data.query_params) fields = external_data.form_data @@ -774,20 +783,20 @@ class HttpClient(Client): """ try: headers = dict_copy(self.headers) - kwargs = {"headers": headers, "timeout": 3, "preload_content": True} + kwargs: dict[str, Any] = {"headers": headers, "timeout": 3, "preload_content": True} if self.server_host_name: kwargs["assert_same_host"] = False headers["Host"] = self.server_host_name - response = self.http.request("GET", f"{self.url}/ping", **kwargs) + response = cast(PoolManager, self.http).request("GET", f"{self.url}/ping", **kwargs) return 200 <= response.status < 300 except HTTPError: logger.debug("ping failed", exc_info=True) return False def close_connections(self) -> None: - self.http.clear() + cast(PoolManager, self.http).clear() def close(self) -> None: if self._owns_pool_manager: - self.http.clear() - all_managers.pop(self.http, None) + cast(PoolManager, self.http).clear() + all_managers.pop(cast(PoolManager, self.http), None) diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/driver/httputil.py b/contrib/python/clickhouse-connect/clickhouse_connect/driver/httputil.py index d135e94194d..0f7a6cfa091 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/driver/httputil.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/driver/httputil.py @@ -26,7 +26,7 @@ logger = logging.getLogger(__name__) urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) # Increase this number just to be safe when ClickHouse is returning progress headers -http.client._MAXHEADERS = 10000 +http.client._MAXHEADERS = 10000 # type: ignore[attr-defined] DEFAULT_KEEP_INTERVAL = 30 DEFAULT_KEEP_COUNT = 3 @@ -42,8 +42,8 @@ core_socket_options = [ ] logging.getLogger("urllib3").setLevel(logging.WARNING) -_proxy_managers = {} -all_managers = {} +_proxy_managers: dict[str, PoolManager] = {} +all_managers: dict[PoolManager, int] = {} @atexit.register @@ -74,7 +74,7 @@ def get_pool_manager_options( if getattr(socket, "TCP_KEEPCNT", None) is not None: socket_options.append((SOCKET_TCP, socket.TCP_KEEPCNT, keep_count)) if getattr(socket, "TCP_KEEPIDLE", None) is not None: - socket_options.append((SOCKET_TCP, socket.TCP_KEEPIDLE, keep_idle)) + socket_options.append((SOCKET_TCP, socket.TCP_KEEPIDLE, keep_idle)) # type: ignore[attr-defined] if sys.platform == "darwin": socket_options.append((SOCKET_TCP, getattr(socket, "TCP_KEEPALIVE", 0x10), keep_interval)) options["maxsize"] = options.get("maxsize", 8) @@ -119,7 +119,7 @@ def get_pool_manager( raise ProgrammingError("Only one of http_proxy or https_proxy should be specified") if not http_proxy.startswith("http"): http_proxy = f"http://{http_proxy}" - manager = ProxyManager(http_proxy, **options) + manager: PoolManager = ProxyManager(http_proxy, **options) elif https_proxy: if not https_proxy.startswith("http"): https_proxy = f"https://{https_proxy}" @@ -238,7 +238,6 @@ class ResponseSource: done = False current_size = 0 read_gen = response.stream(chunk_size, decompress is None) - data_received = False read_error = None while True: while not done: @@ -246,7 +245,7 @@ class ResponseSource: try: chunk = next(read_gen, None) # Always try to read at least one chunk if there are any left except Exception as ex: - # Store the exception for potential re-raising if no data was received + # Store the exception for re-raising later read_error = ex logger.warning("unexpected failure to read next chunk", exc_info=True) if not chunk: @@ -257,7 +256,7 @@ class ResponseSource: if current_size > buffer_size: break if len(chunks) == 0: - if read_error and not data_received: + if read_error: raise OperationalError("Failed to read response data from server") from read_error return if decompress: @@ -267,7 +266,6 @@ class ResponseSource: chunk = chunks.popleft() current_size -= len(chunk) if chunk: - data_received = True yield chunk self.gen = buffered() diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/driver/insert.py b/contrib/python/clickhouse-connect/clickhouse_connect/driver/insert.py index bcb361afaac..200063fc940 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/driver/insert.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/driver/insert.py @@ -97,6 +97,7 @@ class InsertContext(BaseQueryContext): self.block_row_count = self._calc_block_size() def _calc_block_size(self) -> int: + assert self._data is not None if self.req_block_size: return self.req_block_size row_size = 0 diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/driver/npquery.py b/contrib/python/clickhouse-connect/clickhouse_connect/driver/npquery.py index 1c68ed71566..01388f1fb8a 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/driver/npquery.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/driver/npquery.py @@ -1,6 +1,7 @@ import itertools import logging from collections.abc import Generator, Sequence +from typing import Any from clickhouse_connect.driver import options from clickhouse_connect.driver.common import StreamContext, empty_gen @@ -13,19 +14,19 @@ logger = logging.getLogger(__name__) class NumpyResult(Closable): def __init__( self, - block_gen: Generator[Sequence, None, None] = None, + block_gen: Generator[Sequence, None, None] | None = None, column_names: tuple = (), column_types: tuple = (), d_types: Sequence = (), - source: Closable = None, + source: Closable | None = None, ): self.column_names = column_names self.column_types = column_types self.np_types = d_types self.source = source self.query_id = "" - self.summary = {} - self._block_gen = block_gen or empty_gen() + self.summary: dict[str, Any] = {} + self._block_gen: Generator[Sequence, None, None] | None = block_gen or empty_gen() self._numpy_result = None self._df_result = None diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/driver/parser.py b/contrib/python/clickhouse-connect/clickhouse_connect/driver/parser.py index 99de3f56cc6..3afe4904fcd 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/driver/parser.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/driver/parser.py @@ -74,7 +74,7 @@ def parse_callable(expr) -> tuple[str, tuple[str | int, ...], str]: return name, tuple(values), expr[pos:].strip() -def parse_enum(expr) -> tuple[tuple[str], tuple[int]]: +def parse_enum(expr) -> tuple[tuple[str, ...], tuple[int, ...]]: """ Parse a ClickHouse enum definition expression of the form ('key1' = 1, 'key2' = 2) :param expr: ClickHouse enum expression/arguments @@ -84,8 +84,8 @@ def parse_enum(expr) -> tuple[tuple[str], tuple[int]]: values = [] pos = expr.find("(") + 1 in_key = False - key = [] - value = [] + key: list[str] = [] + value: list[str] = [] while True: char = expr[pos] pos += 1 @@ -110,8 +110,8 @@ def parse_enum(expr) -> tuple[tuple[str], tuple[int]]: in_key = True else: value.append(char) - values, keys = zip(*sorted(zip(values, keys))) - return tuple(keys), tuple(values) + sorted_values, sorted_keys = zip(*sorted(zip(values, keys))) + return tuple(sorted_keys), tuple(sorted_values) def parse_columns(expr: str): diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/driver/query.py b/contrib/python/clickhouse-connect/clickhouse_connect/driver/query.py index 9338e6ff572..f4a94be4cba 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/driver/query.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/driver/query.py @@ -52,7 +52,7 @@ class QueryContext(BaseQueryContext): def __init__( self, query: str | bytes = "", - parameters: dict[str, Any] | None = None, + parameters: Sequence | dict[str, Any] | None = None, settings: dict[str, Any] | None = None, query_formats: dict[str, str] | None = None, column_formats: dict[str, str | dict[str, str]] | None = None, @@ -117,7 +117,7 @@ class QueryContext(BaseQueryContext): self.parameters = parameters or {} self.use_none = True if use_none is None else use_none self.column_oriented = False if column_oriented is None else column_oriented - self.use_numpy = use_numpy + self.use_numpy = use_numpy if use_numpy is not None else False self.max_str_len = 0 if max_str_len is None else max_str_len self.server_tz = server_tz self.apply_server_tz = apply_server_tz @@ -132,7 +132,7 @@ class QueryContext(BaseQueryContext): raise ProgrammingError(f"query_tz {query_tz} is not recognized; {tzutil.TZDATA_HINT}") from ex self.query_tz = query_tz if column_tzs is not None: - resolved_column_tzs = {} + resolved_column_tzs: dict[str, str | tzinfo] = {} for col_name, col_tz in column_tzs.items(): if isinstance(col_tz, str): try: @@ -143,8 +143,8 @@ class QueryContext(BaseQueryContext): resolved_column_tzs[col_name] = col_tz column_tzs = resolved_column_tzs self.column_tzs = column_tzs - self.column_tz = None - self.response_tz = None + self.column_tz: str | tzinfo | None = None + self.response_tz: tzinfo | None = None self.block_info = False self.as_pandas = as_pandas self.streaming = streaming @@ -177,12 +177,12 @@ class QueryContext(BaseQueryContext): def is_command(self) -> bool: return command_re.search(self.uncommented_query) is not None or bare_row_policy_show_re.search(self.uncommented_query) is not None - def set_parameters(self, parameters: dict[str, Any]): + def set_parameters(self, parameters: Sequence | dict[str, Any]): self.parameters = parameters self._update_query() def set_parameter(self, key: str, value: Any): - if not self.parameters: + if not isinstance(self.parameters, dict): self.parameters = {} self.parameters[key] = value self._update_query() @@ -219,7 +219,7 @@ class QueryContext(BaseQueryContext): def updated_copy( self, query: str | bytes | None = None, - parameters: dict[str, Any] | None = None, + parameters: Sequence | dict[str, Any] | None = None, settings: dict[str, Any] | None = None, query_formats: dict[str, str] | None = None, column_formats: dict[str, str | dict[str, str]] | None = None, @@ -245,7 +245,11 @@ class QueryContext(BaseQueryContext): resolved_tz_mode = tz_mode if tz_mode is not None else self.tz_mode return QueryContext( query=query or self.query, - parameters=dict_copy(self.parameters, parameters), + parameters=( + dict_copy(self.parameters, parameters if isinstance(parameters, dict) else None) + if isinstance(self.parameters, dict) + else (parameters if parameters is not None else self.parameters) + ), settings=dict_copy(self.settings, settings), query_formats=dict_copy(self.query_formats, query_formats), column_formats=dict_copy(self.column_formats, column_formats), @@ -283,18 +287,18 @@ class QueryResult(Closable): def __init__( self, - result_set: Matrix = None, - block_gen: Generator[Matrix, None, None] = None, + result_set: Matrix | None = None, + block_gen: Generator[Matrix, None, None] | None = None, column_names: tuple[str, ...] = (), column_types: tuple["ClickHouseType", ...] = (), column_oriented: bool = False, - source: Closable = None, - query_id: str = None, - summary: dict[str, Any] = None, + source: Closable | None = None, + query_id: str | None = None, + summary: dict[str, Any] | None = None, ): - self._result_rows = result_set - self._result_columns = None - self._block_gen = block_gen or empty_gen() + self._result_rows: Matrix | None = result_set + self._result_columns: Matrix | None = None + self._block_gen: Generator[Matrix, None, None] | None = block_gen or empty_gen() self._in_context = False self._query_id = query_id self.column_names = column_names @@ -320,7 +324,7 @@ class QueryResult(Closable): else: self._result_columns = [[] for _ in range(len(self.column_names))] else: - result = [[] for _ in range(len(self.column_names))] + result: list[list[Any]] = [[] for _ in range(len(self.column_names))] with self.column_block_stream as stream: for block in stream: for base, added in zip(result, block): @@ -343,7 +347,7 @@ class QueryResult(Closable): query_id = self.summary.get("query_id") if query_id: return query_id - return self._query_id + return self._query_id or "" def _column_block_stream(self): if self._block_gen is None: @@ -383,13 +387,17 @@ class QueryResult(Closable): return len(self.result_set) @property - def first_item(self) -> dict[str, Any]: + def first_item(self) -> dict[str, Any] | None: + if self.row_count == 0: + return None if self.column_oriented: return {name: col[0] for name, col in zip(self.column_names, self.result_set)} return dict(zip(self.column_names, self.result_set[0])) @property - def first_row(self) -> Sequence[Any]: + def first_row(self) -> Sequence[Any] | None: + if self.row_count == 0: + return None if self.column_oriented: return [col[0] for col in self.result_set] return self.result_set[0] diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/driver/streaming.py b/contrib/python/clickhouse-connect/clickhouse_connect/driver/streaming.py index 0a20bfb7697..ddd50f21b3e 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/driver/streaming.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/driver/streaming.py @@ -33,15 +33,15 @@ class StreamingResponseSource(Closable): self.exception_tag = exception_tag # maxsize=10 means max ~10 socket reads buffered - self.queue = AsyncSyncQueue(maxsize=10) + self.queue: AsyncSyncQueue[bytes | Exception] = AsyncSyncQueue(maxsize=10) self._decompressor = None self._decompressor_initialized = False # Multiple accesses to .gen must return the same generator, not create new ones - self._gen_cache = None + self._gen_cache: Iterator[bytes] | None = None - self._producer_task = None + self._producer_task: asyncio.Task | None = None self._producer_started = threading.Event() self._producer_error: Exception | None = None self._producer_completed = False @@ -257,7 +257,7 @@ class StreamingInsertSource: self.transform = transform self.context = context self.loop = loop - self.queue = AsyncSyncQueue(maxsize=maxsize) + self.queue: AsyncSyncQueue[bytes | bytearray | Exception] = AsyncSyncQueue(maxsize=maxsize) self._producer_future = None self._started = False diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/driver/summary.py b/contrib/python/clickhouse-connect/clickhouse_connect/driver/summary.py index 8938bb82ed1..6627bc65ad0 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/driver/summary.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/driver/summary.py @@ -3,9 +3,9 @@ from clickhouse_connect.driver.query import QueryResult class QuerySummary: - summary = {} + summary: dict[str, str] = {} - def __init__(self, summary: dict | None = None): + def __init__(self, summary: dict[str, str] | None = None): if summary is not None: self.summary = summary @@ -20,7 +20,7 @@ class QuerySummary: return self.summary.get("query_id", "") def as_query_result(self) -> QueryResult: - data = [] + data: list[int | str] = [] column_names = [] column_types = [] str_type = get_from_name("String") diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/driver/transform.py b/contrib/python/clickhouse-connect/clickhouse_connect/driver/transform.py index 5970277b5ee..4a4378d239e 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/driver/transform.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/driver/transform.py @@ -3,7 +3,7 @@ import logging from clickhouse_connect.datatypes import registry from clickhouse_connect.driver.common import write_leb128 from clickhouse_connect.driver.compression import get_compressor -from clickhouse_connect.driver.exceptions import StreamCompleteException, StreamFailureError +from clickhouse_connect.driver.exceptions import OperationalError, StreamCompleteException, StreamFailureError from clickhouse_connect.driver.insert import InsertContext from clickhouse_connect.driver.npquery import NumpyResult from clickhouse_connect.driver.query import QueryContext, QueryResult @@ -71,8 +71,10 @@ class NativeTransform: raise StreamFailureError(error_msg) from None raise StreamFailureError("Stream ended unexpectedly (connection closed by server)") from ex - # Handle async streaming errors (ClientPayloadError from aiohttp) - if ex.__class__.__name__ == "ClientPayloadError": + # A read failure partway through the stream: OperationalError from the sync reader, + # ClientPayloadError from aiohttp. ClickHouse may have written the real error into the + # response body before the connection dropped, so prefer that over the transport error. + if isinstance(ex, OperationalError) or ex.__class__.__name__ == "ClientPayloadError": if source.last_message: error_msg = None exception_tag = getattr(source, "exception_tag", None) @@ -106,7 +108,8 @@ class NativeTransform: @staticmethod def build_insert(context: InsertContext): - compressor = get_compressor(context.compression) + compression = context.compression if isinstance(context.compression, str) else None + compressor = get_compressor(compression) def chunk_gen(): for block in context.next_block(): diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/driver/types.py b/contrib/python/clickhouse-connect/clickhouse_connect/driver/types.py index f54b5e86e60..7dadbac32bf 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/driver/types.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/driver/types.py @@ -12,7 +12,7 @@ class Closable(ABC): class ByteSource(Closable): - last_message: bytes = None + last_message: bytes | None = None @abstractmethod def read_leb128(self) -> int: @@ -31,7 +31,7 @@ class ByteSource(Closable): pass @abstractmethod - def read_str_col(self, num_rows: int, encoding: str, nullable: bool = False, null_obj: Any = None): + def read_str_col(self, num_rows: int, encoding: str | None, nullable: bool = False, null_obj: Any = None): pass @abstractmethod diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/driver/tzutil.py b/contrib/python/clickhouse-connect/clickhouse_connect/driver/tzutil.py index 7eb330da8bb..f00065df60d 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/driver/tzutil.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/driver/tzutil.py @@ -5,7 +5,7 @@ from datetime import datetime, timedelta, timezone, tzinfo tzlocal = None try: - import tzlocal # Maybe we can use the tzlocal module to get a safe timezone + import tzlocal # type: ignore[no-redef] # Maybe we can use the tzlocal module to get a safe timezone except ImportError: pass @@ -270,7 +270,9 @@ def _detect_local_tz() -> tzinfo: return resolve_zone(env_tz) except zoneinfo.ZoneInfoNotFoundError: pass - return datetime.now().astimezone().tzinfo + local = datetime.now().astimezone().tzinfo + assert local is not None + return local local_tz, local_tz_dst_safe = normalize_timezone(_detect_local_tz()) diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/driverc/buffer.pyx b/contrib/python/clickhouse-connect/clickhouse_connect/driverc/buffer.pyx index 1e8bc8e0ca2..9eac5a94d70 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/driverc/buffer.pyx +++ b/contrib/python/clickhouse-connect/clickhouse_connect/driverc/buffer.pyx @@ -1,3 +1,4 @@ +# cython: freethreading_compatible = True import sys from typing import Iterable, Any, Optional @@ -22,12 +23,19 @@ cdef char * errors = 'strict' cdef char * utf8 = 'utf8' cdef dict array_templates = {} cdef bint must_swap = sys.byteorder == 'big' -cdef array.array swapper = array.array('Q', [0]) -for c in 'bBuhHiIlLqQfd': +for c in 'bBhHiIlLqQfd': array_templates[c] = array.array(c, []) +cdef inline unsigned long long _bswap_uint64(unsigned long long v): + """Byte-swap a 64-bit unsigned integer for big-endian systems.""" + return (((v & 0xFF) << 56) | (((v >> 8) & 0xFF) << 48) | + (((v >> 16) & 0xFF) << 40) | (((v >> 24) & 0xFF) << 32) | + (((v >> 32) & 0xFF) << 24) | (((v >> 40) & 0xFF) << 16) | + (((v >> 48) & 0xFF) << 8) | ((v >> 56) & 0xFF)) + + cdef class ResponseBuffer: def __init__(self, source): self.slice_sz = 4096 @@ -279,11 +287,11 @@ cdef class ResponseBuffer: @cython.wraparound(False) def read_uint64(self) -> int: cdef ull_wrapper* x + cdef unsigned long long tmp cdef char* b = self.read_bytes_c(8) if must_swap: - memcpy(swapper.data.as_voidptr, b, 8) - swapper.byteswap() - return swapper[0] + memcpy(&tmp, b, 8) + return _bswap_uint64(tmp) x = <ull_wrapper *> b return x.int_value diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/driverc/dataconv.pyx b/contrib/python/clickhouse-connect/clickhouse_connect/driverc/dataconv.pyx index 2718ae46454..c82fe85b5e3 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/driverc/dataconv.pyx +++ b/contrib/python/clickhouse-connect/clickhouse_connect/driverc/dataconv.pyx @@ -1,3 +1,4 @@ +# cython: freethreading_compatible = True import struct from typing import Sequence, Optional diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/driverc/npconv.pyx b/contrib/python/clickhouse-connect/clickhouse_connect/driverc/npconv.pyx index 21568f4f5de..120db6638f5 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/driverc/npconv.pyx +++ b/contrib/python/clickhouse-connect/clickhouse_connect/driverc/npconv.pyx @@ -1,3 +1,4 @@ +# cython: freethreading_compatible = True import cython import numpy as np diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/py.typed b/contrib/python/clickhouse-connect/clickhouse_connect/py.typed new file mode 100644 index 00000000000..e69de29bb2d --- /dev/null +++ b/contrib/python/clickhouse-connect/clickhouse_connect/py.typed diff --git a/contrib/python/clickhouse-connect/ya.make b/contrib/python/clickhouse-connect/ya.make index 34bf63b5b78..800f19f6818 100644 --- a/contrib/python/clickhouse-connect/ya.make +++ b/contrib/python/clickhouse-connect/ya.make @@ -2,7 +2,7 @@ PY3_LIBRARY() -VERSION(1.3.0) +VERSION(1.4.1) LICENSE(Apache-2.0) @@ -119,6 +119,7 @@ RESOURCE_FILES( .dist-info/METADATA .dist-info/entry_points.txt .dist-info/top_level.txt + clickhouse_connect/py.typed ) END() diff --git a/contrib/python/textual/.dist-info/METADATA b/contrib/python/textual/.dist-info/METADATA index f3bb97ae79a..bdcd6666e66 100644 --- a/contrib/python/textual/.dist-info/METADATA +++ b/contrib/python/textual/.dist-info/METADATA @@ -1,6 +1,6 @@ Metadata-Version: 2.4 Name: textual -Version: 8.2.7 +Version: 8.2.8 Summary: Modern Text User Interface framework License: MIT License-File: LICENSE diff --git a/contrib/python/textual/textual/_compositor.py b/contrib/python/textual/textual/_compositor.py index ee4c231f240..9af38a23c7e 100644 --- a/contrib/python/textual/textual/_compositor.py +++ b/contrib/python/textual/textual/_compositor.py @@ -923,11 +923,8 @@ class Compositor: x -= region.x + gutter_left y -= region.y + gutter_right - if y < 0: - return None, None - - if x < 0: - return widget, Offset(0, y) + if x < 0 or y < 0: + return widget, Offset(max(0, x), max(0, y)) visible_screen_stack.set(widget.app._background_screens) line = widget.render_line(y) diff --git a/contrib/python/textual/textual/_xterm_parser.py b/contrib/python/textual/textual/_xterm_parser.py index eafa7414ecb..df9112370cd 100644 --- a/contrib/python/textual/textual/_xterm_parser.py +++ b/contrib/python/textual/textual/_xterm_parser.py @@ -39,7 +39,7 @@ SPECIAL_SEQUENCES = {BRACKETED_PASTE_START, BRACKETED_PASTE_END, FOCUSIN, FOCUSO """Set of special sequences.""" _re_extended_key: Final[re.Pattern[str]] = re.compile( - r"\x1b\[((?:\d*;?){2,3})([u~ABCDEFHPQRS])" + r"\x1b\[((?:[\d:]*;?){2,3})([u~ABCDEFHPQRS])" ) _re_in_band_window_resize: Final[re.Pattern[str]] = re.compile( r"\x1b\[48;(\d+(?:\:.*?)?);(\d+(?:\:.*?)?);(\d+(?:\:.*?)?);(\d+(?:\:.*?)?)t" @@ -334,8 +334,26 @@ class XTermParser(Parser[Message]): self._debug_log_file.close() self._debug_log_file = None + @classmethod + def _parse_colon_codepoints(cls, text_str: str) -> list[str | None]: + """Convert codepoints split on colons in to a list of characters. + + Args: + text_str: String with groups of digits, separated by one or more colons. + + Returns: + A list of characters. + """ + if not text_str: + return [None] + characters: list[str | None] = [ + chr(int(part)) if part.isdecimal() else chr(1) + for part in text_str.split(":") + ] + return characters + @lru_cache(maxsize=1024) - def _parse_extended_key(self, sequence: str) -> events.Key | None: + def _parse_extended_key(self, sequence: str) -> list[events.Key] | None: """Parse a Kitty sequence. Args: @@ -348,38 +366,47 @@ class XTermParser(Parser[Message]): if (match := _re_extended_key.fullmatch(sequence)) is None: return None + key_events: list[events.Key] = [] + codes, end = match.groups(default="") codepoint_str, modifiers_str, text_str, *_ = codes.split(";") + ["", "", ""] - codepoint = int(codepoint_str or "1") modifiers = int(modifiers_str or "0") - text = chr(int(text_str)) if text_str else None - if not (key := FUNCTIONAL_KEYS.get(f"{codepoint}{end}", "")): - key = _character_to_key(text if text else chr(codepoint)) + for text in self._parse_colon_codepoints(text_str): + if not (key := FUNCTIONAL_KEYS.get(f"{codepoint}{end}", "")): + key = _character_to_key(text if text else chr(codepoint)) - key_tokens: list[str] = [] - # The modifier is redundant on a modifier key - if modifiers and key not in MODIFIER_FUNCTIONAL_KEYS and text_str is not None: - modifier_bits = int(modifiers) - 1 - # Not convinced of the utility in reporting caps_lock and num_lock - MODIFIERS = ("alt", "ctrl", "super", "hyper", "meta") - # Ignore caps_lock and num_lock modifiers - if modifier_bits & 1 and (text is None or text.isspace()): - key_tokens.append("shift") - for bit, modifier in enumerate(MODIFIERS, 1): - if modifier == "alt" and text is not None: - continue - if modifier_bits & (1 << bit): - key_tokens.append(modifier) + key_tokens: list[str] = [] + # The modifier is redundant on a modifier key + if ( + modifiers + and key not in MODIFIER_FUNCTIONAL_KEYS + and text_str is not None + ): + modifier_bits = int(modifiers) - 1 + # Not convinced of the utility in reporting caps_lock and num_lock + MODIFIERS = ("alt", "ctrl", "super", "hyper", "meta") + # Ignore caps_lock and num_lock modifiers + if modifier_bits & 1 and (text is None or text.isspace()): + key_tokens.append("shift") + for bit, modifier in enumerate(MODIFIERS, 1): + if modifier == "alt" and text is not None: + continue + if modifier_bits & (1 << bit): + key_tokens.append(modifier) - key_tokens.sort() - if key is not None: - key_tokens.append(key) - return events.Key( - "+".join(key_tokens), - text or (None if modifiers else SPECIAL_KEY_TO_CHARACTER.get(key, None)), - ) + key_tokens.sort() + if key is not None: + key_tokens.append(key) + key_events.append( + events.Key( + "+".join(key_tokens), + text + or (None if modifiers else SPECIAL_KEY_TO_CHARACTER.get(key, None)), + ) + ) + return key_events def _sequence_to_key_events( self, sequence: str, alt: bool = False @@ -395,9 +422,10 @@ class XTermParser(Parser[Message]): if ( not constants.DISABLE_KITTY_KEY - and (key := self._parse_extended_key(sequence)) is not None + and (keys := self._parse_extended_key(sequence)) is not None ): - yield key.copy() + for key in keys: + yield key.copy() return keys = ANSI_SEQUENCES_KEYS.get(sequence) diff --git a/contrib/python/textual/textual/pilot.py b/contrib/python/textual/textual/pilot.py index c387b258276..76adae9d594 100644 --- a/contrib/python/textual/textual/pilot.py +++ b/contrib/python/textual/textual/pilot.py @@ -438,7 +438,7 @@ class Pilot(Generic[ReturnType]): ) offset = Offset(message_arguments["x"], message_arguments["y"]) - if offset not in screen.region: + if offset not in screen.size.region: raise OutOfBounds( "Target offset is outside of currently-visible screen region." ) diff --git a/contrib/python/textual/textual/screen.py b/contrib/python/textual/textual/screen.py index af9070538e7..c8b1d1462d4 100644 --- a/contrib/python/textual/textual/screen.py +++ b/contrib/python/textual/textual/screen.py @@ -1838,8 +1838,11 @@ class Screen(Generic[ScreenResultType], Widget): if select_offset is not None: content_widget = select_widget content_offset = select_offset - assert isinstance(content_widget.parent, Widget) - container = content_widget.parent + container = ( + content_widget + if isinstance(content_widget, Screen) + else content_widget.parent + ) else: content_widget = None container = select_widget @@ -1884,6 +1887,7 @@ class Screen(Generic[ScreenResultType], Widget): select_widget, select_offset = self.get_widget_and_offset_at( event.x, event.y ) + if ( select_widget is not None and select_widget.allow_select @@ -1893,8 +1897,11 @@ class Screen(Generic[ScreenResultType], Widget): if select_offset is not None: content_widget = select_widget content_offset = select_offset - assert isinstance(content_widget.parent, Widget) - container = content_widget.parent + container = ( + content_widget + if isinstance(content_widget, Screen) + else content_widget.parent + ) else: content_widget = None container = select_widget diff --git a/contrib/python/textual/textual/widgets/_input.py b/contrib/python/textual/textual/widgets/_input.py index 3aa2a945bdb..bcf2c2994de 100644 --- a/contrib/python/textual/textual/widgets/_input.py +++ b/contrib/python/textual/textual/widgets/_input.py @@ -122,9 +122,14 @@ class Input(ScrollView): Binding( "ctrl+w", "delete_left_word", "Delete left to start of word", show=False ), - Binding("ctrl+u", "delete_left_all", "Delete all to the left", show=False), Binding( - "ctrl+backspace", + "ctrl+u,super+backspace", + "delete_left_all", + "Delete all to the left", + show=False, + ), + Binding( + "ctrl+backspace,alt+backspace", "delete_right_word", "Delete right to start of word", show=False, @@ -145,6 +150,7 @@ class Input(ScrollView): | shift+right | Move cursor right and select. | | ctrl+right | Move the cursor one word to the right. | | backspace | Delete the character to the left of the cursor. | + | ctrl+backspace,alt+backspace | Delete right to start of word. | | ctrl+shift+right | Move cursor right a word and select. | | ctrl+shift+a | Select all text in the input. | | home,ctrl+a | Go to the beginning of the input. | @@ -154,7 +160,7 @@ class Input(ScrollView): | delete,ctrl+d | Delete the character to the right of the cursor. | | enter | Submit the current value of the input. | | ctrl+w | Delete the word to the left of the cursor. | - | ctrl+u | Delete everything to the left of the cursor. | + | ctrl+u,super+backspace | Delete everything to the left of the cursor. | | ctrl+f | Delete the word to the right of the cursor. | | ctrl+k | Delete everything to the right of the cursor. | | ctrl+x | Cut selected text. | diff --git a/contrib/python/textual/textual/widgets/_text_area.py b/contrib/python/textual/textual/widgets/_text_area.py index d970a31c5a7..1ca74675e54 100644 --- a/contrib/python/textual/textual/widgets/_text_area.py +++ b/contrib/python/textual/textual/widgets/_text_area.py @@ -354,7 +354,7 @@ TextArea { show=False, ), Binding( - "ctrl+w,ctrl+backspace", + "ctrl+w,ctrl+backspace,alt+backspace", "delete_word_left", "Delete left to start of word", show=False, @@ -390,7 +390,7 @@ TextArea { show=False, ), Binding( - "ctrl+u", + "ctrl+u,super+backspace", "delete_to_start_of_line", "Delete to line start", show=False, @@ -446,7 +446,7 @@ TextArea { | delete,ctrl+d | Delete character to the right of cursor. | | alt+delete | Delete from cursor to end of the word. | | ctrl+shift+k | Delete the current line. | - | ctrl+u | Delete from cursor to the start of the line. | + | ctrl+u,super+backspace | Delete from cursor to the start of the line. | | ctrl+k | Delete from cursor to the end of the line. | | f6 | Select the current line. | | f7 | Select all text in the document. | diff --git a/contrib/python/textual/ya.make b/contrib/python/textual/ya.make index 84a9205d17e..fd182d0171f 100644 --- a/contrib/python/textual/ya.make +++ b/contrib/python/textual/ya.make @@ -2,7 +2,7 @@ PY3_LIBRARY() -VERSION(8.2.7) +VERSION(8.2.8) LICENSE(MIT) diff --git a/library/cpp/yt/error/error-inl.h b/library/cpp/yt/error/error-inl.h index 1542d008f58..85751e8a3cf 100644 --- a/library/cpp/yt/error/error-inl.h +++ b/library/cpp/yt/error/error-inl.h @@ -604,4 +604,24 @@ void FormatValue(TStringBuilderBase* builder, const TErrorOr<T>& error, TStringB //////////////////////////////////////////////////////////////////////////////// +template <class F, class... As> +auto RunNoExcept(F&& functor, As&&... args) noexcept -> decltype(functor(std::forward<As>(args)...)) +{ + return functor(std::forward<As>(args)...); +} + +//////////////////////////////////////////////////////////////////////////////// + +inline TStringBuf GetWellKnownLoggingTag(const std::exception&) +{ + return "Error"_sb; +} + +inline TStringBuf GetWellKnownLoggingTag(const TError&) +{ + return "Error"_sb; +} + +//////////////////////////////////////////////////////////////////////////////// + } // namespace NYT diff --git a/library/cpp/yt/error/error.h b/library/cpp/yt/error/error.h index 13223af8d6c..1e5c2147d82 100644 --- a/library/cpp/yt/error/error.h +++ b/library/cpp/yt/error/error.h @@ -475,10 +475,15 @@ void FormatValue(TStringBuilderBase* builder, const TErrorOr<T>& error, TStringB //////////////////////////////////////////////////////////////////////////////// template <class F, class... As> -auto RunNoExcept(F&& functor, As&&... args) noexcept -> decltype(functor(std::forward<As>(args)...)) -{ - return functor(std::forward<As>(args)...); -} +auto RunNoExcept(F&& functor, As&&... args) noexcept -> decltype(functor(std::forward<As>(args)...)); + +//////////////////////////////////////////////////////////////////////////////// + +//! Registers errors as a well-known logging tag (ADL customization point for +//! library/cpp/yt/logging), so that |YT_TLOG_*(...).With(error)| attaches them under +//! the "Error" key, rendered after the message in plain text. +TStringBuf GetWellKnownLoggingTag(const std::exception&); +TStringBuf GetWellKnownLoggingTag(const TError&); //////////////////////////////////////////////////////////////////////////////// diff --git a/library/cpp/yt/logging/backends/arcadia/backend.cpp b/library/cpp/yt/logging/backends/arcadia/backend.cpp index 3c6ff9f5f58..5878ba52be3 100644 --- a/library/cpp/yt/logging/backends/arcadia/backend.cpp +++ b/library/cpp/yt/logging/backends/arcadia/backend.cpp @@ -57,7 +57,7 @@ public: // Use low-level api, because it is more convinient here. auto loggingContext = GetLoggingContext(); auto event = NDetail::CreateLogEvent(loggingContext, Logger_, logLevel); - event.MessageRef = NDetail::BuildLogMessage(loggingContext, Logger_, message).MessageRef; + event.Payload = NDetail::BuildLogMessage(loggingContext, Logger_, message).Payload; event.Family = ELogFamily::PlainText; Logger_.Write(std::move(event)); } diff --git a/library/cpp/yt/logging/benchmark/logging.cpp b/library/cpp/yt/logging/benchmark/logging.cpp new file mode 100644 index 00000000000..8745511f2bd --- /dev/null +++ b/library/cpp/yt/logging/benchmark/logging.cpp @@ -0,0 +1,167 @@ +#include <benchmark/benchmark.h> + +#include <library/cpp/yt/logging/logger.h> + +#include <atomic> + +namespace NYT::NLogging { +namespace { + +//////////////////////////////////////////////////////////////////////////////// + +// A log manager whose Enqueue is essentially free, so that the benchmark measures +// only the producer-side cost (building the event payload), not the logging thread. +class TBenchmarkLogManager + : public ILogManager +{ +public: + void RegisterStaticAnchor(TLoggingAnchor* anchor, ::TSourceLocation /*sourceLocation*/, TStringBuf /*message*/) override + { + anchor->Registered.store(true); + } + + void UpdateAnchor(TLoggingAnchor* anchor) override + { + anchor->CurrentVersion.store(ActualVersion_.load()); + } + + void Enqueue(TLogEvent&& event) override + { + // Keep the produced payload observable so the build is not optimized away; + // the event is then destroyed (freeing its chunk slice) -- same for both APIs. + benchmark::DoNotOptimize(std::get<TTaggedLogEventPayload>(event.Payload).Underlying().Begin()); + } + + const TLoggingCategory* GetCategory(TStringBuf /*categoryName*/) override + { + return &Category_; + } + + void UpdateCategory(TLoggingCategory* category) override + { + category->CurrentVersion.store(ActualVersion_.load()); + } + + bool GetAbortOnAlert() const override + { + return false; + } + +private: + std::atomic<int> ActualVersion_ = 1; + TLoggingCategory Category_ = { + .Name = "Bench", + .MinPlainTextLevel = ELogLevel::Minimum, + .CurrentVersion = 0, + .ActualVersion = &ActualVersion_, + }; +}; + +//////////////////////////////////////////////////////////////////////////////// +// No tags. + +void BM_Log_NoTags(benchmark::State& state) +{ + TBenchmarkLogManager manager; + TLogger Logger(&manager, "Bench"); + for (auto _ : state) { + YT_LOG_INFO("This is a log message of a typical length without tags"); + } +} +BENCHMARK(BM_Log_NoTags); + +void BM_TLog_NoTags(benchmark::State& state) +{ + TBenchmarkLogManager manager; + TLogger Logger(&manager, "Bench"); + for (auto _ : state) { + YT_TLOG_INFO("This is a log message of a typical length without tags"); + } +} +BENCHMARK(BM_TLog_NoTags); + +//////////////////////////////////////////////////////////////////////////////// +// One int tag. + +void BM_Log_OneTag(benchmark::State& state) +{ + TBenchmarkLogManager manager; + TLogger Logger(&manager, "Bench"); + for (auto _ : state) { + YT_LOG_INFO("This is a log message of a typical length (TagName: %v)", + 123); + } +} +BENCHMARK(BM_Log_OneTag); + +void BM_TLog_OneTag(benchmark::State& state) +{ + TBenchmarkLogManager manager; + TLogger Logger(&manager, "Bench"); + for (auto _ : state) { + YT_TLOG_INFO("This is a log message of a typical length") + .With("TagName", 123); + } +} +BENCHMARK(BM_TLog_OneTag); + +//////////////////////////////////////////////////////////////////////////////// +// Two tags (int + string). + +void BM_Log_TwoTags(benchmark::State& state) +{ + TBenchmarkLogManager manager; + TLogger Logger(&manager, "Bench"); + for (auto _ : state) { + YT_LOG_INFO("This is a log message of a typical length (TagName1: %v, TagName2: %v)", + 123, + "this-is-a-string"); + } +} +BENCHMARK(BM_Log_TwoTags); + +void BM_TLog_TwoTags(benchmark::State& state) +{ + TBenchmarkLogManager manager; + TLogger Logger(&manager, "Bench"); + for (auto _ : state) { + YT_TLOG_INFO("This is a log message of a typical length") + .With("TagName1", 123) + .With("TagName2", "this-is-a-string"); + } +} +BENCHMARK(BM_TLog_TwoTags); + +//////////////////////////////////////////////////////////////////////////////// +// Three tags (int + string + double). + +void BM_Log_ThreeTags(benchmark::State& state) +{ + TBenchmarkLogManager manager; + TLogger Logger(&manager, "Bench"); + for (auto _ : state) { + YT_LOG_INFO("This is a log message of a typical length (TagName1: %v, TagName2: %v, TagName3: %v)", + 123, + "this-is-a-string", + 3.14); + } +} +BENCHMARK(BM_Log_ThreeTags); + +void BM_TLog_ThreeTags(benchmark::State& state) +{ + TBenchmarkLogManager manager; + TLogger Logger(&manager, "Bench"); + for (auto _ : state) { + YT_TLOG_INFO("This is a log message of a typical length") + .With("TagName1", 123) + .With("TagName2", "this-is-a-string") + .With("TagName3", 3.14); + } +} +BENCHMARK(BM_TLog_ThreeTags); + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace +} // namespace NYT::NLogging diff --git a/library/cpp/yt/logging/benchmark/ya.make b/library/cpp/yt/logging/benchmark/ya.make index 61fd0b49276..a7a23128cdf 100644 --- a/library/cpp/yt/logging/benchmark/ya.make +++ b/library/cpp/yt/logging/benchmark/ya.make @@ -4,6 +4,7 @@ INCLUDE(${ARCADIA_ROOT}/library/cpp/yt/ya_cpp.make.inc) SRCS( logger_tag.cpp + logging.cpp ) PEERDIR( diff --git a/library/cpp/yt/logging/logger-inl.h b/library/cpp/yt/logging/logger-inl.h index fa20fb06489..ff4d6d81bbe 100644 --- a/library/cpp/yt/logging/logger-inl.h +++ b/library/cpp/yt/logging/logger-inl.h @@ -4,6 +4,8 @@ #include "logger.h" #endif +#include "tagged_payload.h" + #include <library/cpp/yt/yson_string/convert.h> #include <library/cpp/yt/yson_string/string.h> @@ -107,33 +109,6 @@ Y_FORCE_INLINE const TLogger& TLogger::operator()() const namespace NDetail { -struct TMessageStringBuilderContext -{ - TSharedMutableRef Chunk; -}; - -struct TMessageBufferTag -{ }; - -class TMessageStringBuilder - : public TStringBuilderBase -{ -public: - TSharedRef Flush(); - - // For testing only. - static void DisablePerThreadCache(); - -protected: - void DoReset() override; - void DoReserve(size_t newLength) override; - -private: - TSharedMutableRef Buffer_; - - static constexpr size_t ChunkSize = 128_KB - 64; -}; - inline bool HasMessageTags( const TLoggingContext& loggingContext, const TLogger& logger) @@ -162,14 +137,14 @@ inline void AppendMessageTags( } if (const auto& traceLoggingTag = loggingContext.TraceLoggingTag; !traceLoggingTag.empty()) { if (printComma) { - builder->AppendString(TStringBuf(", ")); + builder->AppendString(", "_sb); } builder->AppendString(traceLoggingTag); printComma = true; } if (const auto& threadMessageTag = GetThreadMessageTag(); !threadMessageTag.empty()) { if (printComma) { - builder->AppendString(TStringBuf(", ")); + builder->AppendString(", "_sb); } builder->AppendString(threadMessageTag); printComma = true; @@ -185,10 +160,10 @@ inline void AppendLogMessage( if (HasMessageTags(loggingContext, logger)) { if (message.Size() >= 1 && message[message.Size() - 1] == ')') { builder->AppendString(TStringBuf(message.Begin(), message.Size() - 1)); - builder->AppendString(TStringBuf(", ")); + builder->AppendString(", "_sb); } else { builder->AppendString(TStringBuf(message.Begin(), message.Size())); - builder->AppendString(TStringBuf(" (")); + builder->AppendString(" ("_sb); } AppendMessageTags(builder, loggingContext, logger); builder->AppendChar(')'); @@ -208,10 +183,10 @@ void AppendLogMessageWithFormat( if (HasMessageTags(loggingContext, logger)) { if (format.size() >= 2 && format[format.size() - 1] == ')') { builder->AppendFormat(TRuntimeFormat{format.substr(0, format.size() - 1)}, std::forward<TArgs>(args)...); - builder->AppendString(TStringBuf(", ")); + builder->AppendString(", "_sb); } else { builder->AppendFormat(TRuntimeFormat{format}, std::forward<TArgs>(args)...); - builder->AppendString(TStringBuf(" (")); + builder->AppendString(" ("_sb); } AppendMessageTags(builder, loggingContext, logger); builder->AppendChar(')'); @@ -222,7 +197,7 @@ void AppendLogMessageWithFormat( struct TLogMessage { - TSharedRef MessageRef; + TTaggedLogEventPayload Payload; TStringBuf Anchor; }; @@ -233,9 +208,10 @@ TLogMessage BuildLogMessage( TFormatString<TArgs...> format, TArgs&&... args) { - TMessageStringBuilder builder; - AppendLogMessageWithFormat(&builder, loggingContext, logger, format.Get(), std::forward<TArgs>(args)...); - return {builder.Flush(), format.Get()}; + TTaggedPayloadWriter writer; + AppendLogMessageWithFormat(writer.BeginMessage(), loggingContext, logger, format.Get(), std::forward<TArgs>(args)...); + writer.EndMessage(); + return {writer.Finish(), format.Get()}; } template <CFormattable T> @@ -245,13 +221,15 @@ TLogMessage BuildLogMessage( const TLogger& logger, const T& obj) { - TMessageStringBuilder builder; - FormatValue(&builder, obj, TStringBuf("v")); + TTaggedPayloadWriter writer; + auto* builder = writer.BeginMessage(); + FormatValue(builder, obj, "v"_sb); if (HasMessageTags(loggingContext, logger)) { - builder.AppendString(TStringBuf(" (")); - AppendMessageTags(&builder, loggingContext, logger); - builder.AppendChar(')'); + builder->AppendString(" ("_sb); + AppendMessageTags(builder, loggingContext, logger); + builder->AppendChar(')'); } + writer.EndMessage(); if constexpr (std::same_as<TStringBuf, std::remove_cvref_t<T>>) { // NB(arkady-e1ppa): This is the overload where TStringBuf @@ -262,9 +240,9 @@ TLogMessage BuildLogMessage( // us having overload for TStringBuf (both have implicit ctors from // string literals) thus we have to accommodate TStringBuf specifics // in this if constexpr part. - return {builder.Flush(), obj}; + return {writer.Finish(), obj}; } else { - return {builder.Flush(), TStringBuf()}; + return {writer.Finish(), TStringBuf()}; } } @@ -295,13 +273,10 @@ inline TLogMessage BuildLogMessage( const TLogger& logger, TSharedRef&& message) { - if (HasMessageTags(loggingContext, logger)) { - TMessageStringBuilder builder; - AppendLogMessage(&builder, loggingContext, logger, message); - return {builder.Flush(), TStringBuf()}; - } else { - return {std::move(message), TStringBuf()}; - } + TTaggedPayloadWriter writer; + AppendLogMessage(writer.BeginMessage(), loggingContext, logger, message); + writer.EndMessage(); + return {writer.Finish(), TStringBuf()}; } inline TLogEvent CreateLogEvent( @@ -310,10 +285,10 @@ inline TLogEvent CreateLogEvent( ELogLevel level) { TLogEvent event; - event.Instant = loggingContext.Instant; event.Category = logger.GetCategory(); - event.Essential = logger.IsEssential(); event.Level = level; + event.Essential = logger.IsEssential(); + event.Instant = loggingContext.Instant; event.ThreadId = loggingContext.ThreadId; event.ThreadName = loggingContext.ThreadName; event.FiberId = loggingContext.FiberId; @@ -332,15 +307,24 @@ inline void LogEventImpl( ELogLevel level, ::TSourceLocation sourceLocation, TLoggingAnchor* anchor, - TSharedRef message) + TTaggedLogEventPayload payload) { - auto event = CreateLogEvent(loggingContext, logger, level); - event.MessageKind = ELogMessageKind::Unstructured; - event.MessageRef = std::move(message); - event.Family = ELogFamily::PlainText; - event.SourceFile = sourceLocation.File; - event.SourceLine = sourceLocation.Line; - event.Anchor = anchor; + auto event = TLogEvent{ + .Category = logger.GetCategory(), + .Level = level, + .Family = ELogFamily::PlainText, + .Essential = logger.IsEssential(), + .Payload = std::move(payload), + .Instant = loggingContext.Instant, + .ThreadId = loggingContext.ThreadId, + .ThreadName = loggingContext.ThreadName, + .FiberId = loggingContext.FiberId, + .TraceId = loggingContext.TraceId, + .RequestId = loggingContext.RequestId, + .SourceFile = sourceLocation.File, + .SourceLine = sourceLocation.Line, + .Anchor = anchor, + }; if (Y_UNLIKELY(event.Level >= ELogLevel::Alert)) { logger.Write(TLogEvent(event)); OnCriticalLogEvent(logger, event); @@ -349,6 +333,299 @@ inline void LogEventImpl( } } +//////////////////////////////////////////////////////////////////////////////// + +//! References the per-call-site static anchor and its one-shot registration flag. +//! Produced by the lambda embedded in the fluent |YT_TLOG_*| macros. +struct TStaticAnchorRef +{ + TLoggingAnchor* Anchor; + std::atomic<bool>* Registered; +}; + +//! Wraps the format spec passed to |TTaggedLoggingGuard::With| and validates at compile +//! time that it is a |%|-prefixed string literal (e.g. |"%v"|, |"%08x"|). The stored +//! spec has the leading |%| stripped, as expected by |FormatValue|. +class TLoggingTagSpec +{ +public: + template <size_t N> + consteval TLoggingTagSpec(const char (&spec)[N]) + : Spec_(spec + 1, N - 2) + { + static_assert(N >= 2, "Logging tag format spec must be a non-empty string literal"); + if (spec[0] != '%') { + TheLoggingTagFormatSpecMustStartWithPercentSign(); + } + } + + TStringBuf Get() const + { + return Spec_; + } + +private: + const TStringBuf Spec_; + + // Undefined on purpose: calling it from the |consteval| ctor turns a missing + // leading |%| into a compile error that names the violated rule. + static void TheLoggingTagFormatSpecMustStartWithPercentSign(); +}; + +class TWellKnownTaggedLoggingGuard; + +//! Accumulates a tagged log message via a fluent |.With| chain and emits the event in +//! its destructor. Instantiated by the fluent |YT_TLOG_*| macros, which guarantee that +//! the chain is reached only when the level is enabled (so tag value expressions are +//! not evaluated otherwise). +//! +//! The user message -- with the logger's contextual (logger/trace/thread) tags folded +//! in -- goes to the payload message field; each |.With(key, value)| becomes a +//! structured payload tag (see #TTaggedPayloadWriter). +class TTaggedLoggingGuard +{ +public: + TTaggedLoggingGuard( + const TLogger& logger, + ELogLevel level, + ::TSourceLocation sourceLocation, + TStaticAnchorRef anchorRef, + TStringBuf message) + : TTaggedLoggingGuard( + logger, + level, + sourceLocation, + anchorRef, + message, + /*alwaysBuildMessage*/ false) + { } + + TTaggedLoggingGuard(const TTaggedLoggingGuard&) = delete; + TTaggedLoggingGuard& operator=(const TTaggedLoggingGuard&) = delete; + + bool IsEnabled() const + { + return Enabled_; + } + + template <class TValue> + TTaggedLoggingGuard& With(TStringBuf tag, const TValue& value) & + { + return DoWith(tag, value, "v"_sb); + } + + template <class TValue> + TTaggedLoggingGuard& With(TStringBuf tag, const TValue& value, TLoggingTagSpec spec) & + { + return DoWith(tag, value, spec.Get()); + } + + //! Attaches a well-known tag whose key is resolved from #value's type via the + //! |GetWellKnownLoggingTag| ADL point (the type must opt in, e.g. errors). + //! + //! Returns a #TWellKnownTaggedLoggingGuard, which exposes only further well-known + //! tags: the payload contract requires well-known tags to come last (so + //! #FormatTaggedPayload can stay single-pass), so a keyed |.With(key, value)| after a + //! well-known tag must not compile. + template <class TValue> + TWellKnownTaggedLoggingGuard With(const TValue& value) &; + + ~TTaggedLoggingGuard() + { + if (!Enabled_) { + return; + } + LogEventImpl( + LoggingContext_, + Logger_, + EffectiveLevel_, + SourceLocation_, + Anchor_, + Writer_.Finish()); + } + +protected: + const TLogger& Logger_; + const ::TSourceLocation SourceLocation_; + TLoggingAnchor* const Anchor_; + + bool Enabled_ = false; + ELogLevel EffectiveLevel_ = ELogLevel::Minimum; + TLoggingContext LoggingContext_; + TTaggedPayloadWriter Writer_; + + //! Shared constructor. When #alwaysBuildMessage is set the payload message is built + //! even if the level is disabled (so a terminal guard can still recover it); #Enabled_ + //! continues to gate whether the destructor emits the event. + TTaggedLoggingGuard( + const TLogger& logger, + ELogLevel level, + ::TSourceLocation sourceLocation, + TStaticAnchorRef anchorRef, + TStringBuf message, + bool alwaysBuildMessage) + : Logger_(logger) + , SourceLocation_(sourceLocation) + , Anchor_(anchorRef.Anchor) + { + if (!Logger_.IsAnchorUpToDate(*Anchor_)) [[unlikely]] { + Logger_.UpdateStaticAnchor(Anchor_, anchorRef.Registered, sourceLocation, message); + } + + EffectiveLevel_ = TLogger::GetEffectiveLoggingLevel(level, *Anchor_); + Enabled_ = Logger_.IsLevelEnabled(EffectiveLevel_); + if (!Enabled_ && !alwaysBuildMessage) { + return; + } + + LoggingContext_ = GetLoggingContext(); + + auto* builder = Writer_.BeginMessage(); + builder->AppendString(message); + if (HasMessageTags(LoggingContext_, Logger_)) { + builder->AppendString(" ("_sb); + AppendMessageTags(builder, LoggingContext_, Logger_); + builder->AppendChar(')'); + } + Writer_.EndMessage(); + } + +private: + template <class TValue> + TTaggedLoggingGuard& DoWith(TStringBuf tag, const TValue& value, TStringBuf spec) & + { + // Format the value straight into the payload buffer; no temporary. + FormatValue(Writer_.BeginTag(tag), value, spec); + Writer_.EndTag(); + return *this; + } +}; + +//! Restricts the fluent |.With| chain once a well-known tag has been attached. The +//! payload contract requires well-known tags to come last (#FormatTaggedPayload is +//! single-pass), so only further well-known |.With(value)| calls are exposed -- a keyed +//! |.With(key, value)| after a well-known tag fails to compile. +class TWellKnownTaggedLoggingGuard +{ +public: + explicit TWellKnownTaggedLoggingGuard(TTaggedLoggingGuard& guard) + : Guard_(guard) + { } + + template <class TValue> + TWellKnownTaggedLoggingGuard With(const TValue& value) && + { + return Guard_.With(value); + } + +private: + TTaggedLoggingGuard& Guard_; +}; + +template <class TValue> +TWellKnownTaggedLoggingGuard TTaggedLoggingGuard::With(const TValue& value) & +{ + FormatValue(Writer_.BeginWellKnownTag(GetWellKnownLoggingTag(value)), value, "v"_sb); + Writer_.EndTag(); + return TWellKnownTaggedLoggingGuard(*this); +} + +//! Terminal guard for the fluent |YT_TLOG_FATAL| macros. Builds the message +//! unconditionally and, once the |.With| chain completes, emits the event at |Fatal| +//! level -- which aborts the process. The enclosing |for| invokes #Commit in its step. +class TTaggedFatalLoggingGuard + : public TTaggedLoggingGuard +{ +public: + TTaggedFatalLoggingGuard( + const TLogger& logger, + ::TSourceLocation sourceLocation, + TStaticAnchorRef anchorRef, + TStringBuf message) + : TTaggedLoggingGuard(logger, ELogLevel::Fatal, sourceLocation, anchorRef, message, /*alwaysBuildMessage*/ true) + { } + + //! Returns true exactly once, so the enclosing |for| runs the |.With| chain a single + //! time before its step expression commits the event. + bool TryEnter() + { + bool pending = Pending_; + Pending_ = false; + return pending; + } + + //! Emits the event at |Fatal| level; the log manager aborts the process. + [[noreturn]] void Commit() & + { + Enabled_ = false; // The event is emitted here, not from the base destructor. + LogEventImpl(LoggingContext_, Logger_, ELogLevel::Fatal, SourceLocation_, Anchor_, Writer_.Finish()); + Y_UNREACHABLE(); + } + +private: + bool Pending_ = true; +}; + +//! Terminal guard for the fluent |YT_TLOG_ALERT_AND_THROW| macros. Builds the message +//! unconditionally; once the |.With| chain completes, #Commit emits the event at |Alert| +//! level (when enabled) and returns the message so the macro can attach it to the thrown +//! error. The throw lives in the macro -- the logging library must not depend on the +//! error library, and a destructor must not throw. +class TTaggedThrowingLoggingGuard + : public TTaggedLoggingGuard +{ +public: + TTaggedThrowingLoggingGuard( + const TLogger& logger, + ::TSourceLocation sourceLocation, + TStaticAnchorRef anchorRef, + TStringBuf message) + : TTaggedLoggingGuard(logger, ELogLevel::Alert, sourceLocation, anchorRef, message, /*alwaysBuildMessage*/ true) + { } + + //! Returns true exactly once, so the enclosing |for| runs the |.With| chain a single + //! time before its step expression commits the event and throws. + bool TryEnter() + { + bool pending = Pending_; + Pending_ = false; + return pending; + } + + //! Emits the alert event (when enabled) and returns its message for the error payload. + std::string Commit() & + { + auto payload = Writer_.Finish(); + std::string message(GetMessageFromTaggedPayload(payload)); + if (Enabled_) { + Enabled_ = false; // The event is emitted here, not from the base destructor. + LogEventImpl(LoggingContext_, Logger_, EffectiveLevel_, SourceLocation_, Anchor_, std::move(payload)); + } + return message; + } + +private: + bool Pending_ = true; +}; + +//! A no-op stand-in for #TTaggedLoggingGuard used by compile-time-disabled trace logging: +//! it swallows the |.With| chain without evaluating it. +class TNullTaggedLoggingGuard +{ +public: + template <class... TArgs> + TNullTaggedLoggingGuard& With(TArgs&&...) + { + return *this; + } +}; + +template <class TMessage> +TNullTaggedLoggingGuard MakeNullTaggedLoggingGuard(const TMessage&) +{ + return {}; +} + } // namespace NDetail //////////////////////////////////////////////////////////////////////////////// diff --git a/library/cpp/yt/logging/logger.cpp b/library/cpp/yt/logging/logger.cpp index acb823a614e..5359c1cd141 100644 --- a/library/cpp/yt/logging/logger.cpp +++ b/library/cpp/yt/logging/logger.cpp @@ -1,5 +1,7 @@ #include "logger.h" +#include "structured_payload.h" + #include <library/cpp/yt/assert/assert.h> #include <library/cpp/yt/cpu_clock/clock.h> @@ -25,98 +27,13 @@ void OnCriticalLogEvent( event.Level == ELogLevel::Alert && logger.GetAbortOnAlert()) { fprintf(stderr, "*** Aborting on critical log event\n"); - fwrite(event.MessageRef.begin(), 1, event.MessageRef.size(), stderr); + auto message = FormatTaggedPayload(std::get<TTaggedLogEventPayload>(event.Payload)); + fwrite(message.data(), 1, message.size(), stderr); fprintf(stderr, "\n"); YT_ABORT(); } } -TSharedRef TMessageStringBuilder::Flush() -{ - return Buffer_.Slice(0, GetLength()); -} - -void TMessageStringBuilder::DoReset() -{ - Buffer_.Reset(); -} - -struct TPerThreadCache; - -YT_DEFINE_THREAD_LOCAL(TPerThreadCache*, Cache); -YT_DEFINE_THREAD_LOCAL(bool, CacheDestroyed); - -struct TPerThreadCache -{ - TSharedMutableRef Chunk; - size_t ChunkOffset = 0; - - ~TPerThreadCache() - { - TMessageStringBuilder::DisablePerThreadCache(); - } - - static YT_PREVENT_TLS_CACHING TPerThreadCache* GetCache() - { - auto& cache = Cache(); - if (Y_LIKELY(cache)) { - return cache; - } - if (CacheDestroyed()) { - return nullptr; - } - static thread_local TPerThreadCache CacheData; - cache = &CacheData; - return cache; - } -}; - -void TMessageStringBuilder::DisablePerThreadCache() -{ - Cache() = nullptr; - CacheDestroyed() = true; -} - -void TMessageStringBuilder::DoReserve(size_t newCapacity) -{ - auto oldLength = GetLength(); - newCapacity = FastClp2(newCapacity); - - auto newChunkSize = std::max(ChunkSize, newCapacity); - // Hold the old buffer until the data is copied. - auto oldBuffer = std::move(Buffer_); - auto* cache = TPerThreadCache::GetCache(); - if (Y_LIKELY(cache)) { - auto oldCapacity = End_ - Begin_; - auto deltaCapacity = newCapacity - oldCapacity; - if (End_ == cache->Chunk.Begin() + cache->ChunkOffset && - cache->ChunkOffset + deltaCapacity <= cache->Chunk.Size()) - { - // Resize inplace. - Buffer_ = cache->Chunk.Slice(cache->ChunkOffset - oldCapacity, cache->ChunkOffset + deltaCapacity); - cache->ChunkOffset += deltaCapacity; - End_ = Begin_ + newCapacity; - return; - } - - if (Y_UNLIKELY(cache->ChunkOffset + newCapacity > cache->Chunk.Size())) { - cache->Chunk = TSharedMutableRef::Allocate<TMessageBufferTag>(newChunkSize, {.InitializeStorage = false}); - cache->ChunkOffset = 0; - } - - Buffer_ = cache->Chunk.Slice(cache->ChunkOffset, cache->ChunkOffset + newCapacity); - cache->ChunkOffset += newCapacity; - } else { - Buffer_ = TSharedMutableRef::Allocate<TMessageBufferTag>(newChunkSize, {.InitializeStorage = false}); - newCapacity = newChunkSize; - } - if (oldLength > 0) { - ::memcpy(Buffer_.Begin(), Begin_, oldLength); - } - Begin_ = Buffer_.Begin(); - End_ = Begin_ + newCapacity; -} - } // namespace NDetail //////////////////////////////////////////////////////////////////////////////// @@ -370,8 +287,7 @@ void LogStructuredEvent( loggingContext, logger, level); - event.MessageKind = ELogMessageKind::Structured; - event.MessageRef = message.ToSharedRef(); + event.Payload = MakeStructuredPayloadFromYson(message); event.Family = ELogFamily::Structured; logger.Write(std::move(event)); } diff --git a/library/cpp/yt/logging/logger.h b/library/cpp/yt/logging/logger.h index d781326e88d..0f29b6f7376 100644 --- a/library/cpp/yt/logging/logger.h +++ b/library/cpp/yt/logging/logger.h @@ -80,11 +80,6 @@ using TRequestId = TGuid; //////////////////////////////////////////////////////////////////////////////// -DEFINE_ENUM(ELogMessageKind, - (Unstructured) - (Structured) -); - struct TLogEvent { const TLoggingCategory* Category = nullptr; @@ -92,8 +87,9 @@ struct TLogEvent ELogFamily Family = ELogFamily::PlainText; bool Essential = false; - ELogMessageKind MessageKind = ELogMessageKind::Unstructured; - TSharedRef MessageRef; + //! The active alternative identifies the event kind (tagged plain-text vs + //! structured) and determines how the bytes are decoded. + TLogEventPayload Payload; TCpuInstant Instant = 0; @@ -327,10 +323,10 @@ void LogStructuredEvent( #define YT_LOG_ALERT_IF(condition, ...) if (condition) YT_LOG_ALERT(__VA_ARGS__) #define YT_LOG_ALERT_UNLESS(condition, ...) if (!(condition)) YT_LOG_ALERT(__VA_ARGS__) -#define YT_LOG_FATAL(...) \ - do { \ +#define YT_LOG_FATAL(...) \ + do { \ YT_LOG_EVENT(Logger, ::NYT::NLogging::ELogLevel::Fatal, __VA_ARGS__); \ - Y_UNREACHABLE(); \ + Y_UNREACHABLE(); \ } while(false) #define YT_LOG_FATAL_IF(condition, ...) if (Y_UNLIKELY(condition)) YT_LOG_FATAL(__VA_ARGS__) #define YT_LOG_FATAL_UNLESS(condition, ...) if (!Y_LIKELY(condition)) YT_LOG_FATAL(__VA_ARGS__) @@ -346,126 +342,233 @@ void LogStructuredEvent( * 4. The exception is raised regardless of the logging level configured for the |Logger| and * and the currently active message suppressions. */ -#define YT_LOG_ALERT_AND_THROW(...) \ - do { \ - /* NOLINTBEGIN(bugprone-reserved-identifier, readability-identifier-naming) */ \ - const auto& logger__ = Logger(); \ - auto level__ = ::NYT::NLogging::ELogLevel::Alert; \ - auto location__ = __LOCATION__; \ - \ - auto loggingContext__ = ::NYT::NLogging::GetLoggingContext(); \ - auto message__ = ::NYT::NLogging::NDetail::BuildLogMessage(loggingContext__, logger__, __VA_ARGS__); \ - auto messageStr__ = ::std::string(message__.MessageRef.ToStringBuf()); \ - \ - static ::NYT::TLeakyStorage<::NYT::NLogging::TLoggingAnchor> anchorStorage__; \ - auto* anchor__ = anchorStorage__.Get(); \ - \ - bool anchorUpToDate__ = logger__.IsAnchorUpToDate(*anchor__); \ - if (!anchorUpToDate__) [[unlikely]] { \ - static std::atomic<bool> anchorRegistered__; \ - logger__.UpdateStaticAnchor(anchor__, &anchorRegistered__, location__, message__.Anchor); \ - } \ - \ - auto effectiveLevel__ = ::NYT::NLogging::TLogger::GetEffectiveLoggingLevel(level__, *anchor__); \ - if (logger__.IsLevelEnabled(effectiveLevel__)) { \ - ::NYT::NLogging::NDetail::LogEventImpl( \ - loggingContext__, \ - logger__, \ - effectiveLevel__, \ - location__, \ - anchor__, \ - std::move(message__.MessageRef)); \ - } \ - \ - THROW_ERROR_EXCEPTION( \ - ::NYT::EErrorCode::Fatal, \ - "Malformed request or incorrect state detected") \ - << ::NYT::TErrorAttribute("message", std::move(messageStr__)); \ - /* NOLINTEND(bugprone-reserved-identifier, readability-identifier-naming) */ \ +#define YT_LOG_ALERT_AND_THROW(...) \ + do { \ + /* NOLINTBEGIN(bugprone-reserved-identifier, readability-identifier-naming) */ \ + const auto& logger__ = Logger(); \ + auto level__ = ::NYT::NLogging::ELogLevel::Alert; \ + auto location__ = __LOCATION__; \ + \ + auto loggingContext__ = ::NYT::NLogging::GetLoggingContext(); \ + auto message__ = ::NYT::NLogging::NDetail::BuildLogMessage(loggingContext__, logger__, __VA_ARGS__); \ + /* Copy the message out before the payload is moved into the log event below. */ \ + auto messageStr__ = ::std::string(::NYT::NLogging::GetMessageFromTaggedPayload(message__.Payload)); \ + \ + static ::NYT::TLeakyStorage<::NYT::NLogging::TLoggingAnchor> anchorStorage__; \ + auto* anchor__ = anchorStorage__.Get(); \ + \ + bool anchorUpToDate__ = logger__.IsAnchorUpToDate(*anchor__); \ + if (!anchorUpToDate__) [[unlikely]] { \ + static std::atomic<bool> anchorRegistered__; \ + logger__.UpdateStaticAnchor(anchor__, &anchorRegistered__, location__, message__.Anchor); \ + } \ + \ + auto effectiveLevel__ = ::NYT::NLogging::TLogger::GetEffectiveLoggingLevel(level__, *anchor__); \ + if (logger__.IsLevelEnabled(effectiveLevel__)) { \ + ::NYT::NLogging::NDetail::LogEventImpl( \ + loggingContext__, \ + logger__, \ + effectiveLevel__, \ + location__, \ + anchor__, \ + std::move(message__.Payload)); \ + } \ + \ + THROW_ERROR_EXCEPTION( \ + ::NYT::EErrorCode::Fatal, \ + "Malformed request or incorrect state detected") \ + << ::NYT::TErrorAttribute("message", std::move(messageStr__)); \ + /* NOLINTEND(bugprone-reserved-identifier, readability-identifier-naming) */ \ } while (false) #define YT_LOG_ALERT_AND_THROW_IF(condition, ...) \ - if (Y_UNLIKELY(condition)) { \ - YT_LOG_ALERT_AND_THROW(__VA_ARGS__); \ - } \ + if (Y_UNLIKELY(condition)) { \ + YT_LOG_ALERT_AND_THROW(__VA_ARGS__); \ + } \ static_assert(true) #define YT_LOG_ALERT_AND_THROW_UNLESS(condition, ...) \ - if (!Y_UNLIKELY(condition)) { \ - YT_LOG_ALERT_AND_THROW(__VA_ARGS__); \ - } \ + if (!Y_UNLIKELY(condition)) { \ + YT_LOG_ALERT_AND_THROW(__VA_ARGS__); \ + } \ static_assert(true) -#define YT_LOG_EVENT(logger, level, ...) \ - do { \ - /* NOLINTBEGIN(bugprone-reserved-identifier, readability-identifier-naming) */ \ - const auto& logger__ = (logger)(); \ - auto level__ = (level); \ - auto location__ = __LOCATION__; \ - static ::NYT::TLeakyStorage<::NYT::NLogging::TLoggingAnchor> anchorStorage__; \ - auto* anchor__ = anchorStorage__.Get(); \ - \ - bool anchorUpToDate__ = logger__.IsAnchorUpToDate(*anchor__); \ - if (anchorUpToDate__) [[likely]] { \ - auto effectiveLevel__ = ::NYT::NLogging::TLogger::GetEffectiveLoggingLevel(level__, *anchor__); \ - if (!logger__.IsLevelEnabled(effectiveLevel__)) { \ - break; \ - } \ - } \ - \ - auto loggingContext__ = ::NYT::NLogging::GetLoggingContext(); \ +#define YT_LOG_EVENT(logger, level, ...) \ + do { \ + /* NOLINTBEGIN(bugprone-reserved-identifier, readability-identifier-naming) */ \ + const auto& logger__ = (logger)(); \ + auto level__ = (level); \ + auto location__ = __LOCATION__; \ + static ::NYT::TLeakyStorage<::NYT::NLogging::TLoggingAnchor> anchorStorage__; \ + auto* anchor__ = anchorStorage__.Get(); \ + \ + bool anchorUpToDate__ = logger__.IsAnchorUpToDate(*anchor__); \ + if (anchorUpToDate__) [[likely]] { \ + auto effectiveLevel__ = ::NYT::NLogging::TLogger::GetEffectiveLoggingLevel(level__, *anchor__); \ + if (!logger__.IsLevelEnabled(effectiveLevel__)) { \ + break; \ + } \ + } \ + \ + auto loggingContext__ = ::NYT::NLogging::GetLoggingContext(); \ auto message__ = ::NYT::NLogging::NDetail::BuildLogMessage(loggingContext__, logger__, __VA_ARGS__); \ - \ - if (!anchorUpToDate__) [[unlikely]] { \ - static std::atomic<bool> anchorRegistered__; \ - logger__.UpdateStaticAnchor(anchor__, &anchorRegistered__, location__, message__.Anchor); \ - } \ - \ - auto effectiveLevel__ = ::NYT::NLogging::TLogger::GetEffectiveLoggingLevel(level__, *anchor__); \ - if (!logger__.IsLevelEnabled(effectiveLevel__)) { \ - break; \ - } \ - \ - ::NYT::NLogging::NDetail::LogEventImpl( \ - loggingContext__, \ - logger__, \ - effectiveLevel__, \ - location__, \ - anchor__, \ - std::move(message__.MessageRef)); \ - /* NOLINTEND(bugprone-reserved-identifier, readability-identifier-naming) */ \ + \ + if (!anchorUpToDate__) [[unlikely]] { \ + static std::atomic<bool> anchorRegistered__; \ + logger__.UpdateStaticAnchor(anchor__, &anchorRegistered__, location__, message__.Anchor); \ + } \ + \ + auto effectiveLevel__ = ::NYT::NLogging::TLogger::GetEffectiveLoggingLevel(level__, *anchor__); \ + if (!logger__.IsLevelEnabled(effectiveLevel__)) { \ + break; \ + } \ + \ + ::NYT::NLogging::NDetail::LogEventImpl( \ + loggingContext__, \ + logger__, \ + effectiveLevel__, \ + location__, \ + anchor__, \ + std::move(message__.Payload)); \ + /* NOLINTEND(bugprone-reserved-identifier, readability-identifier-naming) */ \ } while (false) -#define YT_LOG_EVENT_WITH_DYNAMIC_ANCHOR(logger, level, anchor, ...) \ - do { \ - const auto& logger__ = (logger)(); \ - auto level__ = (level); \ - auto location__ = __LOCATION__; \ - auto* anchor__ = (anchor); \ - \ - bool anchorUpToDate__ = logger__.IsAnchorUpToDate(*anchor__); \ - if (!anchorUpToDate__) [[unlikely]] { \ - logger__.UpdateDynamicAnchor(anchor__); \ - } \ - \ - auto effectiveLevel__ = ::NYT::NLogging::TLogger::GetEffectiveLoggingLevel(level__, *anchor__); \ - if (!logger__.IsLevelEnabled(effectiveLevel__)) { \ - break; \ - } \ - \ - auto loggingContext__ = ::NYT::NLogging::GetLoggingContext(); \ +#define YT_LOG_EVENT_WITH_DYNAMIC_ANCHOR(logger, level, anchor, ...) \ + do { \ + const auto& logger__ = (logger)(); \ + auto level__ = (level); \ + auto location__ = __LOCATION__; \ + auto* anchor__ = (anchor); \ + \ + bool anchorUpToDate__ = logger__.IsAnchorUpToDate(*anchor__); \ + if (!anchorUpToDate__) [[unlikely]] { \ + logger__.UpdateDynamicAnchor(anchor__); \ + } \ + \ + auto effectiveLevel__ = ::NYT::NLogging::TLogger::GetEffectiveLoggingLevel(level__, *anchor__); \ + if (!logger__.IsLevelEnabled(effectiveLevel__)) { \ + break; \ + } \ + \ + auto loggingContext__ = ::NYT::NLogging::GetLoggingContext(); \ auto message__ = ::NYT::NLogging::NDetail::BuildLogMessage(loggingContext__, logger__, __VA_ARGS__); \ - \ - ::NYT::NLogging::NDetail::LogEventImpl( \ - loggingContext__, \ - logger__, \ - effectiveLevel__, \ - location__, \ - anchor__, \ - std::move(message__.MessageRef)); \ + \ + ::NYT::NLogging::NDetail::LogEventImpl( \ + loggingContext__, \ + logger__, \ + effectiveLevel__, \ + location__, \ + anchor__, \ + std::move(message__.Payload)); \ } while (false) //////////////////////////////////////////////////////////////////////////////// +// Tagged logging +// +// Tags are supplied via a fluent |.With(name, value)| (or |.With(name, value, "%spec")|) +// chain; they are carried as structured key/value pairs in the event payload. A single- +// argument |.With(value)| attaches the value under a statically known key resolved by ADL +// (e.g. |.With(error)| under the "Error" key): +// +// YT_TLOG_INFO("Message") +// .With("Key", value) +// .With("Count", count, "%08x") +// .With(error); +// +// If the message is not logged then the |.With| chain is not evaluated, so tag value +// expressions cost nothing. + +//! Yields a #TStaticAnchorRef for the expansion site: a per-call-site leaky anchor and +//! its one-shot registration flag, produced via an immediately-invoked lambda. +#define YT_TLOG_STATIC_ANCHOR_REF() \ + [] { \ + /* NOLINTBEGIN(bugprone-reserved-identifier, readability-identifier-naming) */ \ + static ::NYT::TLeakyStorage<::NYT::NLogging::TLoggingAnchor> anchorStorage__; \ + static std::atomic<bool> anchorRegistered__; \ + return ::NYT::NLogging::NDetail::TStaticAnchorRef{anchorStorage__.Get(), &anchorRegistered__}; \ + /* NOLINTEND(bugprone-reserved-identifier, readability-identifier-naming) */ \ + }() + +#define YT_TLOG_EVENT_FLUENT(logger, level, message) \ + if (::NYT::NLogging::NDetail::TTaggedLoggingGuard loggingGuard__( \ + (logger)(), \ + (level), \ + __LOCATION__, \ + YT_TLOG_STATIC_ANCHOR_REF(), \ + (message)); \ + !loggingGuard__.IsEnabled()) \ + { } else \ + loggingGuard__ + +#ifdef YT_ENABLE_TRACE_LOGGING +#define YT_TLOG_TRACE(message) YT_TLOG_EVENT_FLUENT(Logger, ::NYT::NLogging::ELogLevel::Trace, message) +#define YT_TLOG_TRACE_IF(condition, message) if (condition) YT_TLOG_TRACE(message) +#define YT_TLOG_TRACE_UNLESS(condition, message) if (!(condition)) YT_TLOG_TRACE(message) +#else +#define YT_TLOG_UNUSED(message) if (true) { } else ::NYT::NLogging::NDetail::MakeNullTaggedLoggingGuard(message) +#define YT_TLOG_TRACE(message) YT_TLOG_UNUSED(message) +#define YT_TLOG_TRACE_IF(condition, message) YT_TLOG_UNUSED(message) +#define YT_TLOG_TRACE_UNLESS(condition, message) YT_TLOG_UNUSED(message) +#endif + +#define YT_TLOG_DEBUG(message) YT_TLOG_EVENT_FLUENT(Logger, ::NYT::NLogging::ELogLevel::Debug, message) +#define YT_TLOG_DEBUG_IF(condition, message) if (condition) YT_TLOG_DEBUG(message) +#define YT_TLOG_DEBUG_UNLESS(condition, message) if (!(condition)) YT_TLOG_DEBUG(message) + +#define YT_TLOG_INFO(message) YT_TLOG_EVENT_FLUENT(Logger, ::NYT::NLogging::ELogLevel::Info, message) +#define YT_TLOG_INFO_IF(condition, message) if (condition) YT_TLOG_INFO(message) +#define YT_TLOG_INFO_UNLESS(condition, message) if (!(condition)) YT_TLOG_INFO(message) + +#define YT_TLOG_WARNING(message) YT_TLOG_EVENT_FLUENT(Logger, ::NYT::NLogging::ELogLevel::Warning, message) +#define YT_TLOG_WARNING_IF(condition, message) if (condition) YT_TLOG_WARNING(message) +#define YT_TLOG_WARNING_UNLESS(condition, message) if (!(condition)) YT_TLOG_WARNING(message) + +#define YT_TLOG_ERROR(message) YT_TLOG_EVENT_FLUENT(Logger, ::NYT::NLogging::ELogLevel::Error, message) +#define YT_TLOG_ERROR_IF(condition, message) if (condition) YT_TLOG_ERROR(message) +#define YT_TLOG_ERROR_UNLESS(condition, message) if (!(condition)) YT_TLOG_ERROR(message) + +#define YT_TLOG_ALERT(message) YT_TLOG_EVENT_FLUENT(Logger, ::NYT::NLogging::ELogLevel::Alert, message) +#define YT_TLOG_ALERT_IF(condition, message) if (condition) YT_TLOG_ALERT(message) +#define YT_TLOG_ALERT_UNLESS(condition, message) if (!(condition)) YT_TLOG_ALERT(message) + +// The terminal action of |YT_TLOG_FATAL|/|YT_TLOG_ALERT_AND_THROW| (aborting or throwing) +// must run *after* the |.With| chain, which the user appends to the macro. A guard cannot +// do it in its destructor -- that would run mid-chain on a temporary, and a throwing +// destructor terminates. So both expand to a single-iteration |for| whose step expression +// fires once the chain (the loop body) has completed. + +#define YT_TLOG_FATAL(message) \ + for (::NYT::NLogging::NDetail::TTaggedFatalLoggingGuard loggingGuard__( \ + Logger(), \ + __LOCATION__, \ + YT_TLOG_STATIC_ANCHOR_REF(), \ + (message)); \ + loggingGuard__.TryEnter(); \ + loggingGuard__.Commit()) \ + loggingGuard__ +#define YT_TLOG_FATAL_IF(condition, message) if (condition) [[unlikely]] YT_TLOG_FATAL(message) +#define YT_TLOG_FATAL_UNLESS(condition, message) if (!(condition)) [[unlikely]] YT_TLOG_FATAL(message) + +// See #YT_LOG_ALERT_AND_THROW for the rationale. The throw lives here -- not in the guard +// -- because the logging library must not depend on the error library. The guard's +// |Commit| logs the alert (when enabled) and returns the message for the |"message"| +// attribute. +#define YT_TLOG_ALERT_AND_THROW(message) \ + for (::NYT::NLogging::NDetail::TTaggedThrowingLoggingGuard loggingGuard__( \ + Logger(), \ + __LOCATION__, \ + YT_TLOG_STATIC_ANCHOR_REF(), \ + (message)); \ + loggingGuard__.TryEnter(); \ + THROW_ERROR_EXCEPTION( \ + ::NYT::EErrorCode::Fatal, \ + "Malformed request or incorrect state detected") \ + << ::NYT::TErrorAttribute("message", loggingGuard__.Commit())) \ + loggingGuard__ +#define YT_TLOG_ALERT_AND_THROW_IF(condition, message) if (condition) [[unlikely]] YT_TLOG_ALERT_AND_THROW(message) +#define YT_TLOG_ALERT_AND_THROW_UNLESS(condition, message) if (!(condition)) [[unlikely]] YT_TLOG_ALERT_AND_THROW(message) + +//////////////////////////////////////////////////////////////////////////////// } // namespace NYT::NLogging diff --git a/library/cpp/yt/logging/plain_text_formatter/formatter.cpp b/library/cpp/yt/logging/plain_text_formatter/formatter.cpp index ab2181113dd..85f0889fa9b 100644 --- a/library/cpp/yt/logging/plain_text_formatter/formatter.cpp +++ b/library/cpp/yt/logging/plain_text_formatter/formatter.cpp @@ -1,9 +1,13 @@ #include "formatter.h" +#include <library/cpp/yt/logging/structured_payload.h> + #include <library/cpp/yt/cpu_clock/clock.h> #include <library/cpp/yt/misc/port.h> +#include <variant> + #ifdef YT_USE_SSE42 #include <emmintrin.h> #include <pmmintrin.h> @@ -146,6 +150,37 @@ void FormatMessage(TBaseFormatter* out, TStringBuf message) } } +// Formats |Message (Key: Value, ...)|, with well-known tags (e.g. an error) appended +// after the |(...)| group. Well-known tags are always written last, so a single pass +// suffices. Every piece -- message, tag keys/values, and the newline separating a +// well-known tag -- goes through FormatMessage and is escaped, so the rendered payload +// stays on a single physical line (a newline is emitted as the literal "\n"). +void FormatPayload(TBaseFormatter* out, const TTaggedLogEventPayload& payload) +{ + TTaggedPayloadReader reader(payload); + FormatMessage(out, reader.ReadMessage()); + bool parenOpen = false; + while (auto tag = reader.TryReadTag()) { + if (tag->IsWellKnown) { + if (parenOpen) { + out->AppendChar(')'); + parenOpen = false; + } + FormatMessage(out, "\n"_sb); + FormatMessage(out, tag->Value); + } else { + out->AppendString(parenOpen ? ", "_sb : " ("_sb); + parenOpen = true; + FormatMessage(out, tag->Key); + out->AppendString(": "_sb); + FormatMessage(out, tag->Value); + } + } + if (parenOpen) { + out->AppendChar(')'); + } +} + //////////////////////////////////////////////////////////////////////////////// void TCachingDateFormatter::Format(TBaseFormatter* buffer, TInstant dateTime, bool printMicroseconds) @@ -186,7 +221,13 @@ void TPlainTextEventFormatter::Format(TBaseFormatter* buffer, const TLogEvent& e buffer->AppendChar('\t'); - FormatMessage(buffer, event.MessageRef.ToStringBuf()); + if (const auto* tagged = std::get_if<TTaggedLogEventPayload>(&event.Payload)) { + FormatPayload(buffer, *tagged); + } else { + // A structured event routed to a plain-text writer: emit its raw YSON fragment + // (escaped, so the record stays a single physical line). + FormatMessage(buffer, GetYsonFromStructuredPayload(std::get<TStructuredLogEventPayload>(event.Payload)).AsStringBuf()); + } buffer->AppendChar('\t'); diff --git a/library/cpp/yt/logging/private.h b/library/cpp/yt/logging/private.h new file mode 100644 index 00000000000..b80f0488048 --- /dev/null +++ b/library/cpp/yt/logging/private.h @@ -0,0 +1,15 @@ +#pragma once + +#include "public.h" + +namespace NYT::NLogging::NDetail { + +//////////////////////////////////////////////////////////////////////////////// + +//! Ref-counted memory tag for payload buffers; used for memory accounting. +struct TMessageBufferTag +{ }; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace NYT::NLogging::NDetail diff --git a/library/cpp/yt/logging/public.h b/library/cpp/yt/logging/public.h index 1e2b59ca0d3..17f8998f1af 100644 --- a/library/cpp/yt/logging/public.h +++ b/library/cpp/yt/logging/public.h @@ -1,6 +1,11 @@ #pragma once #include <library/cpp/yt/misc/enum.h> +#include <library/cpp/yt/misc/strong_typedef.h> + +#include <library/cpp/yt/memory/ref.h> + +#include <variant> namespace NYT::NLogging { @@ -36,4 +41,18 @@ struct ILogManager; //////////////////////////////////////////////////////////////////////////////// +//! Opaque payload of a plain-text log event: a message plus optional key/value tags, +//! framed by #TTaggedPayloadWriter. Produced by the YT_LOG_*/YT_TLOG_* macros. +YT_DEFINE_STRONG_TYPEDEF(TTaggedLogEventPayload, TSharedRef); + +//! Opaque payload of a structured log event: a raw YSON map fragment. +//! Produced by #LogStructuredEvent. +YT_DEFINE_STRONG_TYPEDEF(TStructuredLogEventPayload, TSharedRef); + +//! The payload carried by a #TLogEvent: exactly one of the encodings above. The active +//! alternative both identifies the event kind and determines how the payload is decoded. +using TLogEventPayload = std::variant<TTaggedLogEventPayload, TStructuredLogEventPayload>; + +//////////////////////////////////////////////////////////////////////////////// + } // namespace NYT::NLogging diff --git a/library/cpp/yt/logging/structured_payload.cpp b/library/cpp/yt/logging/structured_payload.cpp new file mode 100644 index 00000000000..1a6ce5ea370 --- /dev/null +++ b/library/cpp/yt/logging/structured_payload.cpp @@ -0,0 +1,19 @@ +#include "structured_payload.h" + +namespace NYT::NLogging { + +//////////////////////////////////////////////////////////////////////////////// + +TStructuredLogEventPayload MakeStructuredPayloadFromYson(const NYson::TYsonString& message) +{ + return TStructuredLogEventPayload(message.ToSharedRef()); +} + +NYson::TYsonStringBuf GetYsonFromStructuredPayload(const TStructuredLogEventPayload& payload) +{ + return NYson::TYsonStringBuf(payload.Underlying().ToStringBuf(), NYson::EYsonType::MapFragment); +} + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace NYT::NLogging diff --git a/library/cpp/yt/logging/structured_payload.h b/library/cpp/yt/logging/structured_payload.h new file mode 100644 index 00000000000..81464f4aa92 --- /dev/null +++ b/library/cpp/yt/logging/structured_payload.h @@ -0,0 +1,25 @@ +#pragma once + +#include "public.h" + +#include <library/cpp/yt/memory/ref.h> + +#include <library/cpp/yt/yson_string/string.h> + +namespace NYT::NLogging { + +//////////////////////////////////////////////////////////////////////////////// + +// Structured events carry the raw YSON map fragment as their opaque payload, with no +// framing. These helpers isolate that representation from callers. + +//! Producer: wraps #message into a payload (zero-copy). +TStructuredLogEventPayload MakeStructuredPayloadFromYson(const NYson::TYsonString& message); + +//! Consumer: views the payload as the YSON map fragment it carries. The result views +//! into #payload, which must outlive it. +NYson::TYsonStringBuf GetYsonFromStructuredPayload(const TStructuredLogEventPayload& payload); + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace NYT::NLogging diff --git a/library/cpp/yt/logging/tagged_payload-inl.h b/library/cpp/yt/logging/tagged_payload-inl.h new file mode 100644 index 00000000000..4b71a5bebef --- /dev/null +++ b/library/cpp/yt/logging/tagged_payload-inl.h @@ -0,0 +1,108 @@ +#ifndef TAGGED_PAYLOAD_INL_H_ +#error "Direct inclusion of this file is not allowed, include tagged_payload.h" +// For the sake of sane code completion. +#include "tagged_payload.h" +#endif + +#include <library/cpp/yt/assert/assert.h> + +#include <cstring> + +namespace NYT::NLogging { + +//////////////////////////////////////////////////////////////////////////////// + +inline TSharedRef TTaggedPayloadBuilder::Flush() +{ + return Buffer_.Slice(0, GetLength()); +} + +template <class T> +void TTaggedPayloadBuilder::AppendPod(const T& value) +{ + ::memcpy(Preallocate(sizeof(value)), &value, sizeof(value)); + Advance(sizeof(value)); +} + +template <class T> +void TTaggedPayloadBuilder::WritePodAt(size_t offset, const T& value) +{ + YT_ASSERT(offset + sizeof(value) <= GetLength()); + ::memcpy(Begin_ + offset, &value, sizeof(value)); +} + +//////////////////////////////////////////////////////////////////////////////// + +inline TStringBuilderBase* TTaggedPayloadWriter::BeginMessage() & +{ + YT_ASSERT(!MessageStarted_); + MessageStarted_ = true; + ReserveLengthPrefix(); + return &Builder_; +} + +inline TTaggedPayloadWriter& TTaggedPayloadWriter::EndMessage() & +{ + YT_ASSERT(MessageStarted_ && !MessageEnded_); + MessageEnded_ = true; + BackpatchLengthPrefix(); + return *this; +} + +inline TStringBuilderBase* TTaggedPayloadWriter::BeginTag(TStringBuf key) & +{ + return DoBeginTag(key, /*wellKnown*/ false); +} + +inline TStringBuilderBase* TTaggedPayloadWriter::BeginWellKnownTag(TStringBuf key) & +{ + return DoBeginTag(key, /*wellKnown*/ true); +} + +inline TStringBuilderBase* TTaggedPayloadWriter::DoBeginTag(TStringBuf key, bool wellKnown) +{ + YT_ASSERT(MessageEnded_ && !InTag_); + // High bit reserved for the well-known flag. + YT_ASSERT(key.size() < WellKnownTagFlag); + InTag_ = true; + // Reserve the key-size prefix, the key bytes, and the value-size prefix in a single + // capacity check; the value is then formatted in place and EndTag backpatches its size. + auto keySize = static_cast<ui32>(key.size()) | (wellKnown ? WellKnownTagFlag : 0); + char* ptr = Builder_.Preallocate(2 * sizeof(ui32) + key.size()); + PrefixOffset_ = Builder_.GetLength() + sizeof(ui32) + key.size(); + ::memcpy(ptr, &keySize, sizeof(keySize)); + ::memcpy(ptr + sizeof(keySize), key.data(), key.size()); + // The value-size prefix is left uninitialized; EndTag overwrites it. + Builder_.Advance(2 * sizeof(ui32) + key.size()); + return &Builder_; +} + +inline TTaggedPayloadWriter& TTaggedPayloadWriter::EndTag() & +{ + YT_ASSERT(InTag_); + InTag_ = false; + BackpatchLengthPrefix(); + return *this; +} + +inline TTaggedLogEventPayload TTaggedPayloadWriter::Finish() & +{ + YT_ASSERT(MessageEnded_ && !InTag_); + return TTaggedLogEventPayload(Builder_.Flush()); +} + +inline void TTaggedPayloadWriter::ReserveLengthPrefix() +{ + PrefixOffset_ = Builder_.GetLength(); + // Reserve a zeroed ui32; backpatched by BackpatchLengthPrefix. + Builder_.AppendPod<ui32>(0); +} + +inline void TTaggedPayloadWriter::BackpatchLengthPrefix() +{ + Builder_.WritePodAt<ui32>(PrefixOffset_, Builder_.GetLength() - PrefixOffset_ - sizeof(ui32)); +} + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace NYT::NLogging diff --git a/library/cpp/yt/logging/tagged_payload.cpp b/library/cpp/yt/logging/tagged_payload.cpp new file mode 100644 index 00000000000..4133cc0bd55 --- /dev/null +++ b/library/cpp/yt/logging/tagged_payload.cpp @@ -0,0 +1,214 @@ +#include "tagged_payload.h" +#include "private.h" + +#include <library/cpp/yt/assert/assert.h> + +#include <library/cpp/yt/misc/tls.h> + +#include <util/system/compiler.h> +#include <util/system/types.h> + +#include <util/generic/bitops.h> + +#include <algorithm> + +namespace NYT::NLogging { + +//////////////////////////////////////////////////////////////////////////////// + +namespace { + +struct TPerThreadCache; + +YT_DEFINE_THREAD_LOCAL(TPerThreadCache*, Cache); +YT_DEFINE_THREAD_LOCAL(bool, CacheDestroyed); + +struct TPerThreadCache +{ + TSharedMutableRef Chunk; + size_t ChunkOffset = 0; + + ~TPerThreadCache() + { + TTaggedPayloadBuilder::DisablePerThreadCache(); + } + + static YT_PREVENT_TLS_CACHING TPerThreadCache* GetCache() + { + auto& cache = Cache(); + [[likely]] if (cache) { + return cache; + } + if (CacheDestroyed()) { + return nullptr; + } + static thread_local TPerThreadCache CacheData; + cache = &CacheData; + return cache; + } +}; + +TSharedMutableRef AllocateChunk(size_t size) +{ + return TSharedMutableRef::Allocate<::NYT::NLogging::NDetail::TMessageBufferTag>(size, {.InitializeStorage = false}); +} + +} // namespace + +//////////////////////////////////////////////////////////////////////////////// + +void TTaggedPayloadBuilder::DisablePerThreadCache() +{ + Cache() = nullptr; + CacheDestroyed() = true; +} + +void TTaggedPayloadBuilder::DoReset() +{ + Buffer_.Reset(); +} + +void TTaggedPayloadBuilder::DoReserve(size_t newCapacity) +{ + auto oldLength = GetLength(); + newCapacity = FastClp2(newCapacity); + + auto newChunkSize = std::max(ChunkSize, newCapacity); + // Hold the old buffer until the data is copied. + auto oldBuffer = std::move(Buffer_); + auto* cache = TPerThreadCache::GetCache(); + [[likely]] if (cache) { + auto oldCapacity = End_ - Begin_; + auto deltaCapacity = newCapacity - oldCapacity; + if (End_ == cache->Chunk.Begin() + cache->ChunkOffset && + cache->ChunkOffset + deltaCapacity <= cache->Chunk.Size()) + { + // Resize inplace. + Buffer_ = cache->Chunk.Slice(cache->ChunkOffset - oldCapacity, cache->ChunkOffset + deltaCapacity); + cache->ChunkOffset += deltaCapacity; + End_ = Begin_ + newCapacity; + return; + } + + [[unlikely]] if (cache->ChunkOffset + newCapacity > cache->Chunk.Size()) { + cache->Chunk = AllocateChunk(newChunkSize); + cache->ChunkOffset = 0; + } + + Buffer_ = cache->Chunk.Slice(cache->ChunkOffset, cache->ChunkOffset + newCapacity); + cache->ChunkOffset += newCapacity; + } else { + Buffer_ = AllocateChunk(newChunkSize); + newCapacity = newChunkSize; + } + if (oldLength > 0) { + ::memcpy(Buffer_.Begin(), Begin_, oldLength); + } + Begin_ = Buffer_.Begin(); + End_ = Begin_ + newCapacity; +} + +//////////////////////////////////////////////////////////////////////////////// + +void TTaggedPayloadWriter::DisablePerThreadCache() +{ + TTaggedPayloadBuilder::DisablePerThreadCache(); +} + +//////////////////////////////////////////////////////////////////////////////// + +TTaggedPayloadReader::TTaggedPayloadReader(const TTaggedLogEventPayload& payload) + : Current_(payload.Underlying().Begin()) + , End_(payload.Underlying().End()) +{ } + +TStringBuf TTaggedPayloadReader::ReadMessage() +{ + YT_ASSERT(!MessageRead_); + MessageRead_ = true; + return ReadString(); +} + +std::optional<TTaggedPayloadReader::TTag> TTaggedPayloadReader::TryReadTag() +{ + YT_ASSERT(MessageRead_); + if (Current_ == End_) { + return std::nullopt; + } + auto keySize = ReadLength(); + bool wellKnown = (keySize & WellKnownTagFlag) != 0; + auto key = ReadBytes(keySize & ~WellKnownTagFlag); + auto value = ReadString(); + return TTag{.IsWellKnown = wellKnown, .Key = key, .Value = value}; +} + +ui32 TTaggedPayloadReader::ReadLength() +{ + YT_ASSERT(Current_ + sizeof(ui32) <= End_); + ui32 size; + ::memcpy(&size, Current_, sizeof(size)); + Current_ += sizeof(size); + return size; +} + +TStringBuf TTaggedPayloadReader::ReadBytes(ui32 size) +{ + YT_ASSERT(Current_ + size <= End_); + TStringBuf result(Current_, size); + Current_ += size; + return result; +} + +TStringBuf TTaggedPayloadReader::ReadString() +{ + return ReadBytes(ReadLength()); +} + +//////////////////////////////////////////////////////////////////////////////// + +TTaggedLogEventPayload MakeTaggedPayloadFromMessage(TStringBuf message) +{ + TTaggedPayloadWriter writer; + writer.BeginMessage()->AppendString(message); + writer.EndMessage(); + return writer.Finish(); +} + +TStringBuf GetMessageFromTaggedPayload(const TTaggedLogEventPayload& payload) +{ + return TTaggedPayloadReader(payload).ReadMessage(); +} + +std::string FormatTaggedPayload(const TTaggedLogEventPayload& payload) +{ + TTaggedPayloadReader reader(payload); + std::string result(reader.ReadMessage()); + // Well-known tags (e.g. an error) are always written last -- the fluent |.With| chain + // enforces this at the type level (see #TWellKnownTaggedLoggingGuard) -- so they render + // after the |(...)| group on trailing lines in a single pass. + bool parenOpen = false; + while (auto tag = reader.TryReadTag()) { + if (tag->IsWellKnown) { + if (parenOpen) { + result += ')'; + parenOpen = false; + } + result += '\n'; + result += tag->Value; + } else { + result += parenOpen ? ", " : " ("; + parenOpen = true; + result += tag->Key; + result += ": "; + result += tag->Value; + } + } + if (parenOpen) { + result += ')'; + } + return result; +} + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace NYT::NLogging diff --git a/library/cpp/yt/logging/tagged_payload.h b/library/cpp/yt/logging/tagged_payload.h new file mode 100644 index 00000000000..da46f30f42a --- /dev/null +++ b/library/cpp/yt/logging/tagged_payload.h @@ -0,0 +1,201 @@ +#pragma once + +#include "public.h" + +#include <library/cpp/yt/string/string_builder.h> + +#include <library/cpp/yt/memory/ref.h> + +#include <util/generic/strbuf.h> +#include <util/generic/size_literals.h> + +#include <util/system/types.h> + +#include <optional> +#include <string> + +namespace NYT::NLogging { + +//////////////////////////////////////////////////////////////////////////////// + +// Plain-text events carry their message together with structured key/value tags in a +// flat, length-prefixed payload, deferring tag rendering to the consumer (the logging +// thread). These helpers own the wire format and isolate it from callers. + +//! A per-thread chunk-cached string builder backing payload buffers. +/*! + * Generic: it knows nothing about the payload layout (see #TTaggedPayloadWriter). + */ +class TTaggedPayloadBuilder + : public TStringBuilderBase +{ +public: + TSharedRef Flush(); + + //! Appends a trivially-copyable value as raw bytes (the ui32 length prefix). + template <class T> + void AppendPod(const T& value); + + //! Overwrites a trivially-copyable value at #offset. + //! #offset + sizeof(value) must not exceed the current length. + template <class T> + void WritePodAt(size_t offset, const T& value); + + //! For testing only. + static void DisablePerThreadCache(); + +private: + TSharedMutableRef Buffer_; + + static constexpr size_t ChunkSize = 128_KB - 64; + + void DoReset() override; + void DoReserve(size_t newCapacity) override; +}; + +//////////////////////////////////////////////////////////////////////////////// + +//! A log message together with its tag key/value pairs, serialized into a flat, +//! opaque byte payload that #TLogEvent hands to the logging thread. +/*! + * Layout (all sizes are host-endian |ui32|; the transport is in-process): + * \code + * [message size][message bytes] + * [key size][key bytes][value size][value bytes] x N + * \endcode + * + * A "well-known" tag sets the high bit of its key-size field; the low 31 bits stay the + * key length and the key bytes are still written. Consumers may render such tags + * specially (see #GetWellKnownLoggingTag). + * + * Both the producer (#TTaggedPayloadWriter) and the consumer + * (#TTaggedPayloadReader) own this single definition of the layout. + */ + +//! The high bit of a tag's key-size field, marking a well-known tag. +constexpr ui32 WellKnownTagFlag = 1u << 31; + +//////////////////////////////////////////////////////////////////////////////// + +//! Producer side: serializes the message and its tags into a payload. +/*! + * Owns a per-thread chunk-cached buffer; both the message text (#BeginMessage) and + * tag values (#BeginTag) are appended through the returned #TStringBuilderBase, so a + * formatted value is built directly into the payload without an intermediate + * allocation/copy. The length-prefix framing lives entirely here -- the underlying + * builder is format-agnostic. + */ +class TTaggedPayloadWriter +{ +public: + //! Constructs an empty writer; the message is built incrementally through the + //! builder returned by #BeginMessage. + TTaggedPayloadWriter() = default; + + TTaggedPayloadWriter(const TTaggedPayloadWriter&) = delete; + TTaggedPayloadWriter& operator=(const TTaggedPayloadWriter&) = delete; + + //! Begins the message field and returns a builder for appending its text. + //! Must be called exactly once, first. + TStringBuilderBase* BeginMessage() &; + + //! Ends the message field, filling in its reserved length prefix. Must be + //! called exactly once, after the message text and before any tag/#Finish. + TTaggedPayloadWriter& EndMessage() &; + + //! Begins a tag: writes #key, then reserves the value's length prefix and returns + //! a builder for appending the value. Lets the value be formatted directly into the + //! payload buffer. Pair with #EndTag. Must follow #EndMessage. + TStringBuilderBase* BeginTag(TStringBuf key) &; + + //! Like #BeginTag, but marks the tag well-known (sets #WellKnownTagFlag). + TStringBuilderBase* BeginWellKnownTag(TStringBuf key) &; + + //! Ends the current tag, filling in the value's reserved length prefix. + TTaggedPayloadWriter& EndTag() &; + + //! Returns the serialized payload. Must follow #EndMessage. + TTaggedLogEventPayload Finish() &; + + //! For testing only. + static void DisablePerThreadCache(); + +private: + TTaggedPayloadBuilder Builder_; + bool MessageStarted_ = false; + bool MessageEnded_ = false; + bool InTag_ = false; + //! Offset of the length prefix currently being filled (message or tag value). + size_t PrefixOffset_ = 0; + + //! Shared implementation of #BeginTag and #BeginWellKnownTag. + TStringBuilderBase* DoBeginTag(TStringBuf key, bool wellKnown); + //! Reserves a length prefix at the current position, remembering its offset. + void ReserveLengthPrefix(); + //! Backpatches the reserved prefix with the length appended after it. + void BackpatchLengthPrefix(); +}; + +//////////////////////////////////////////////////////////////////////////////// + +//! Consumer side: parses a payload produced by #TTaggedPayloadWriter. +/*! + * A single-pass cursor; the returned views point into #payload, which must outlive + * the reader. + * + * Call order: #ReadMessage must be called exactly once, first; then #TryReadTag + * must be called repeatedly until it returns |std::nullopt|, yielding the tags in + * write order. + */ +class TTaggedPayloadReader +{ +public: + struct TTag + { + bool IsWellKnown = false; + TStringBuf Key; + TStringBuf Value; + }; + + explicit TTaggedPayloadReader(const TTaggedLogEventPayload& payload); + + //! Reads the message field. Must be called exactly once, before any #TryReadTag. + TStringBuf ReadMessage(); + + //! Reads the next tag; returns |std::nullopt| once the tags are exhausted. + //! Must follow #ReadMessage. + std::optional<TTag> TryReadTag(); + +private: + const char* Current_; + const char* const End_; + bool MessageRead_ = false; + + //! Reads a ui32 length word. + ui32 ReadLength(); + //! Reads #size bytes as a view. + TStringBuf ReadBytes(ui32 size); + //! Reads a length-prefixed string (a length word followed by its bytes). + TStringBuf ReadString(); +}; + +//////////////////////////////////////////////////////////////////////////////// + +//! Producer convenience: builds a payload carrying only #message and no tags. +TTaggedLogEventPayload MakeTaggedPayloadFromMessage(TStringBuf message); + +//! Consumer convenience: returns the message, disregarding any tags. The result +//! views into #payload, which must outlive it. +TStringBuf GetMessageFromTaggedPayload(const TTaggedLogEventPayload& payload); + +//! Consumer convenience: renders the payload as |Message (Key: Value, ...)|, or just +//! the message if it carries no tags. +std::string FormatTaggedPayload(const TTaggedLogEventPayload& payload); + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace NYT::NLogging + +#define TAGGED_PAYLOAD_INL_H_ +#include "tagged_payload-inl.h" +#undef TAGGED_PAYLOAD_INL_H_ diff --git a/library/cpp/yt/logging/unittests/helpers.cpp b/library/cpp/yt/logging/unittests/helpers.cpp new file mode 100644 index 00000000000..2a2b20120cc --- /dev/null +++ b/library/cpp/yt/logging/unittests/helpers.cpp @@ -0,0 +1,27 @@ +#include "helpers.h" + +namespace NYT::NLogging { + +//////////////////////////////////////////////////////////////////////////////// + +void WriteMessage(TTaggedPayloadWriter* writer, TStringBuf message) +{ + writer->BeginMessage()->AppendString(message); + writer->EndMessage(); +} + +void WriteTag(TTaggedPayloadWriter* writer, TStringBuf key, TStringBuf value) +{ + writer->BeginTag(key)->AppendString(value); + writer->EndTag(); +} + +void WriteWellKnownTag(TTaggedPayloadWriter* writer, TStringBuf key, TStringBuf value) +{ + writer->BeginWellKnownTag(key)->AppendString(value); + writer->EndTag(); +} + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace NYT::NLogging diff --git a/library/cpp/yt/logging/unittests/helpers.h b/library/cpp/yt/logging/unittests/helpers.h new file mode 100644 index 00000000000..6785a65f99b --- /dev/null +++ b/library/cpp/yt/logging/unittests/helpers.h @@ -0,0 +1,22 @@ +#pragma once + +#include <library/cpp/yt/logging/tagged_payload.h> + +#include <util/generic/strbuf.h> + +namespace NYT::NLogging { + +//////////////////////////////////////////////////////////////////////////////// + +//! Writes the message (equivalent to BeginMessage, an append, and EndMessage). +void WriteMessage(TTaggedPayloadWriter* writer, TStringBuf message); + +//! Writes a tag with a ready string value (equivalent to BeginTag, an append, and EndTag). +void WriteTag(TTaggedPayloadWriter* writer, TStringBuf key, TStringBuf value); + +//! Writes a well-known tag with a ready string value (see BeginWellKnownTag). +void WriteWellKnownTag(TTaggedPayloadWriter* writer, TStringBuf key, TStringBuf value); + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace NYT::NLogging diff --git a/library/cpp/yt/logging/unittests/logger_ut.cpp b/library/cpp/yt/logging/unittests/logger_ut.cpp index 410a56fc309..db176ca5394 100644 --- a/library/cpp/yt/logging/unittests/logger_ut.cpp +++ b/library/cpp/yt/logging/unittests/logger_ut.cpp @@ -1,15 +1,117 @@ #include <library/cpp/testing/gtest/gtest.h> #include <library/cpp/yt/logging/logger.h> +#include <library/cpp/yt/logging/tagged_payload.h> +#include <library/cpp/yt/logging/structured_payload.h> #include <library/cpp/yt/error/error.h> +#include <library/cpp/yt/yson_string/string.h> + +#include <atomic> +#include <string> +#include <utility> +#include <vector> + namespace NYT::NLogging { namespace { +using namespace NDetail; + +//////////////////////////////////////////////////////////////////////////////// + +//! A log manager that captures enqueued events for inspection in tests. +class TMockLogManager + : public ILogManager +{ +public: + explicit TMockLogManager(ELogLevel minLevel = ELogLevel::Minimum) + { + Category_.MinPlainTextLevel.store(minLevel); + } + + void RegisterStaticAnchor(TLoggingAnchor* anchor, ::TSourceLocation /*sourceLocation*/, TStringBuf /*message*/) override + { + anchor->Registered.store(true); + } + + void UpdateAnchor(TLoggingAnchor* anchor) override + { + anchor->CurrentVersion.store(ActualVersion_.load()); + } + + void Enqueue(TLogEvent&& event) override + { + Events_.push_back(std::move(event)); + } + + const TLoggingCategory* GetCategory(TStringBuf /*categoryName*/) override + { + return &Category_; + } + + void UpdateCategory(TLoggingCategory* category) override + { + category->CurrentVersion.store(ActualVersion_.load()); + } + + bool GetAbortOnAlert() const override + { + return false; + } + + const std::vector<TLogEvent>& GetEvents() const + { + return Events_; + } + +private: + std::vector<TLogEvent> Events_; + std::atomic<int> ActualVersion_ = 1; + TLoggingCategory Category_ = { + .Name = "Test", + .MinPlainTextLevel = ELogLevel::Minimum, + .CurrentVersion = 0, + .ActualVersion = &ActualVersion_, + }; +}; + +struct TDecodedEvent +{ + std::string Message; + std::vector<std::pair<std::string, std::string>> Tags; +}; + +TDecodedEvent DecodeEvent(const TLogEvent& event) +{ + TTaggedPayloadReader reader(std::get<TTaggedLogEventPayload>(event.Payload)); + TDecodedEvent result; + result.Message = reader.ReadMessage(); + while (auto tag = reader.TryReadTag()) { + result.Tags.emplace_back(tag->Key, tag->Value); + } + return result; +} + +TDecodedEvent DecodeSingleEvent(const TMockLogManager& manager) +{ + EXPECT_EQ(manager.GetEvents().size(), 1u); + return DecodeEvent(manager.GetEvents()[0]); +} + +NYson::TYsonString MakeMapFragment(TStringBuf yson) +{ + return NYson::TYsonString(yson, NYson::EYsonType::MapFragment); +} + +TStringBuf GetStructuredYson(const TLogEvent& event) +{ + return GetYsonFromStructuredPayload(std::get<TStructuredLogEventPayload>(event.Payload)).AsStringBuf(); +} + //////////////////////////////////////////////////////////////////////////////// -TEST(TLogger, NullByDefault) +TEST(TLoggerTest, NullByDefault) { { TLogger logger; @@ -23,7 +125,7 @@ TEST(TLogger, NullByDefault) } } -TEST(TLogger, CopyOfNullLogger) +TEST(TLoggerTest, CopyOfNullLogger) { TLogger nullLogger{/*logManager*/ nullptr, "Category"}; ASSERT_FALSE(nullLogger); @@ -34,7 +136,7 @@ TEST(TLogger, CopyOfNullLogger) EXPECT_FALSE(logger.IsLevelEnabled(ELogLevel::Fatal)); } -TEST(TLogger, LogAlertAndThrowMessage) +TEST(TLoggerTest, LogAlertAndThrowMessage) { try { TLogger Logger; @@ -52,5 +154,228 @@ TEST(TLogger, LogAlertAndThrowMessage) //////////////////////////////////////////////////////////////////////////////// +TEST(TTaggedApiTest, MessageOnly) +{ + TMockLogManager manager; + TLogger Logger(&manager, "Test"); + YT_TLOG_INFO("Message"); + + auto decoded = DecodeSingleEvent(manager); + EXPECT_EQ(decoded.Message, "Message"); + EXPECT_TRUE(decoded.Tags.empty()); +} + +TEST(TTaggedApiTest, Tags) +{ + TMockLogManager manager; + TLogger Logger(&manager, "Test"); + YT_TLOG_INFO("Message") + .With("Arg1", 123) + .With("Arg2", "test"); + + auto decoded = DecodeSingleEvent(manager); + EXPECT_EQ(decoded.Message, "Message"); + ASSERT_EQ(decoded.Tags.size(), 2u); + EXPECT_EQ(decoded.Tags[0], std::pair(std::string("Arg1"), std::string("123"))); + EXPECT_EQ(decoded.Tags[1], std::pair(std::string("Arg2"), std::string("test"))); +} + +TEST(TTaggedApiTest, CustomSpec) +{ + TMockLogManager manager; + TLogger Logger(&manager, "Test"); + YT_TLOG_INFO("Message") + .With("Arg1", 256, "%x"); + + auto decoded = DecodeSingleEvent(manager); + ASSERT_EQ(decoded.Tags.size(), 1u); + EXPECT_EQ(decoded.Tags[0], std::pair(std::string("Arg1"), std::string("100"))); +} + +TEST(TTaggedApiTest, LoggerTagFoldedIntoMessage) +{ + TMockLogManager manager; + auto Logger = TLogger(&manager, "Test") + .WithTag("LoggingTag: %v", 555); + YT_TLOG_INFO("Message") + .With("Arg1", 123); + + auto decoded = DecodeSingleEvent(manager); + EXPECT_EQ(decoded.Message, "Message (LoggingTag: 555)"); + ASSERT_EQ(decoded.Tags.size(), 1u); + EXPECT_EQ(decoded.Tags[0], std::pair(std::string("Arg1"), std::string("123"))); +} + +TEST(TTaggedApiTest, DisabledDoesNotEvaluateTags) +{ + TMockLogManager manager(/*minLevel*/ ELogLevel::Warning); + TLogger Logger(&manager, "Test"); + + int evaluated = 0; + auto evaluate = [&] { + ++evaluated; + return 123; + }; + YT_TLOG_INFO("Message") + .With("Arg1", evaluate()); + + EXPECT_EQ(evaluated, 0); + EXPECT_TRUE(manager.GetEvents().empty()); +} + +TEST(TTaggedApiTest, WellKnownErrorTag) +{ + TMockLogManager manager; + TLogger Logger(&manager, "Test"); + auto error = TError("boom"); + YT_TLOG_INFO("Message") + .With("Tag", 1) + .With(error); + + ASSERT_EQ(manager.GetEvents().size(), 1u); + TTaggedPayloadReader reader(std::get<TTaggedLogEventPayload>(manager.GetEvents()[0].Payload)); + EXPECT_EQ(reader.ReadMessage(), "Message"); + + auto regular = reader.TryReadTag(); + ASSERT_TRUE(regular); + EXPECT_EQ(regular->Key, "Tag"); + EXPECT_EQ(regular->Value, "1"); + EXPECT_FALSE(regular->IsWellKnown); + + // |.With(error)| attaches the error under the well-known "Error" key (resolved via + // GetWellKnownLoggingTag), with the formatted error as the value. + auto errorTag = reader.TryReadTag(); + ASSERT_TRUE(errorTag); + EXPECT_EQ(errorTag->Key, "Error"); + EXPECT_TRUE(errorTag->IsWellKnown); + EXPECT_NE(errorTag->Value.find("boom"), TStringBuf::npos); + + EXPECT_FALSE(reader.TryReadTag()); +} + +// The single-pass formatter assumes well-known tags come last, so the fluent API must +// forbid a keyed tag after a well-known one (e.g. |.With(error).With("Key", 1)|) at +// compile time. +template <class TGuard> +concept CAllowsKeyedTagAfterWellKnown = requires (TGuard guard, TError error) { + guard.With(error).With("Key", 1); +}; + +// A further well-known tag after a well-known one stays allowed. +template <class TGuard> +concept CAllowsWellKnownTagAfterWellKnown = requires (TGuard guard, TError error) { + guard.With(error).With(error); +}; + +static_assert(!CAllowsKeyedTagAfterWellKnown<TTaggedLoggingGuard>); +static_assert(CAllowsWellKnownTagAfterWellKnown<TTaggedLoggingGuard>); + +TEST(TTaggedApiTest, AlertAndThrow) +{ + TMockLogManager manager; + TLogger Logger(&manager, "Test"); + + try { + YT_TLOG_ALERT_AND_THROW("Alert message") + .With("Arg1", 1) + .With("Arg2", 2); + EXPECT_TRUE(false); + } catch (const TErrorException& ex) { + const auto& error = ex.Error(); + EXPECT_EQ(error.GetCode(), NYT::EErrorCode::Fatal); + EXPECT_EQ(error.GetMessage(), "Malformed request or incorrect state detected"); + EXPECT_EQ(error.Attributes().Get<std::string>("message"), "Alert message"); + } + + // The alert is also logged, carrying the structured tags. + auto decoded = DecodeSingleEvent(manager); + EXPECT_EQ(decoded.Message, "Alert message"); + ASSERT_EQ(decoded.Tags.size(), 2u); + EXPECT_EQ(decoded.Tags[0], std::pair(std::string("Arg1"), std::string("1"))); + EXPECT_EQ(decoded.Tags[1], std::pair(std::string("Arg2"), std::string("2"))); +} + +TEST(TTaggedApiTest, AlertAndThrowDisabledStillThrows) +{ + TMockLogManager manager(/*minLevel*/ ELogLevel::Fatal); + TLogger Logger(&manager, "Test"); + + // The level is disabled, so nothing is logged, but the exception (with its message) + // is still raised. + EXPECT_THROW( + YT_TLOG_ALERT_AND_THROW("Alert message") + .With("Arg1", 1), + TErrorException); + EXPECT_TRUE(manager.GetEvents().empty()); +} + +TEST(TTaggedApiDeathTest, Fatal) +{ + EXPECT_DEATH({ + TMockLogManager manager; + TLogger Logger(&manager, "Test"); + YT_TLOG_FATAL("Fatal message") + .With("Arg1", 1); + }, "Fatal message \\(Arg1: 1\\)"); +} + +//////////////////////////////////////////////////////////////////////////////// + +TEST(TStructuredApiTest, StructuredEvent) +{ + TMockLogManager manager; + TLogger Logger(&manager, "Test"); + + LogStructuredEvent(Logger, MakeMapFragment("\"key\"=\"value\""), ELogLevel::Info); + + ASSERT_EQ(manager.GetEvents().size(), 1u); + const auto& event = manager.GetEvents()[0]; + EXPECT_TRUE(std::holds_alternative<TStructuredLogEventPayload>(event.Payload)); + EXPECT_EQ(event.Family, ELogFamily::Structured); + EXPECT_EQ(event.Level, ELogLevel::Info); + // The payload is the raw YSON map fragment. + EXPECT_EQ(GetStructuredYson(event), "\"key\"=\"value\""); +} + +TEST(TStructuredApiTest, PreservesLevel) +{ + TMockLogManager manager; + TLogger Logger(&manager, "Test"); + + LogStructuredEvent(Logger, MakeMapFragment("\"a\"=1;\"b\"=2"), ELogLevel::Warning); + + ASSERT_EQ(manager.GetEvents().size(), 1u); + EXPECT_EQ(manager.GetEvents()[0].Level, ELogLevel::Warning); + EXPECT_EQ(GetStructuredYson(manager.GetEvents()[0]), "\"a\"=1;\"b\"=2"); +} + +TEST(TStructuredApiTest, EmptyFragment) +{ + TMockLogManager manager; + TLogger Logger(&manager, "Test"); + + LogStructuredEvent(Logger, MakeMapFragment(""), ELogLevel::Info); + + ASSERT_EQ(manager.GetEvents().size(), 1u); + EXPECT_EQ(GetStructuredYson(manager.GetEvents()[0]), ""); +} + +TEST(TStructuredApiTest, MultipleEvents) +{ + TMockLogManager manager; + TLogger Logger(&manager, "Test"); + + LogStructuredEvent(Logger, MakeMapFragment("\"i\"=0"), ELogLevel::Debug); + LogStructuredEvent(Logger, MakeMapFragment("\"i\"=1"), ELogLevel::Info); + + ASSERT_EQ(manager.GetEvents().size(), 2u); + EXPECT_EQ(manager.GetEvents()[0].Level, ELogLevel::Debug); + EXPECT_EQ(GetStructuredYson(manager.GetEvents()[0]), "\"i\"=0"); + EXPECT_EQ(manager.GetEvents()[1].Level, ELogLevel::Info); + EXPECT_EQ(GetStructuredYson(manager.GetEvents()[1]), "\"i\"=1"); +} + +//////////////////////////////////////////////////////////////////////////////// + } // namespace } // namespace NYT::NLogging diff --git a/library/cpp/yt/logging/unittests/structured_payload_ut.cpp b/library/cpp/yt/logging/unittests/structured_payload_ut.cpp new file mode 100644 index 00000000000..f15970dcb23 --- /dev/null +++ b/library/cpp/yt/logging/unittests/structured_payload_ut.cpp @@ -0,0 +1,36 @@ +#include <library/cpp/testing/gtest/gtest.h> + +#include <library/cpp/yt/logging/structured_payload.h> + +#include <library/cpp/yt/yson_string/string.h> + +namespace NYT::NLogging { +namespace { + +using namespace NYson; + +//////////////////////////////////////////////////////////////////////////////// + +TYsonString MakeMapFragment(TStringBuf yson) +{ + return TYsonString(yson, EYsonType::MapFragment); +} + +TEST(TStructuredPayloadTest, RoundTrip) +{ + auto payload = MakeStructuredPayloadFromYson(MakeMapFragment("\"key\"=\"value\"")); + auto view = GetYsonFromStructuredPayload(payload); + EXPECT_EQ(view.GetType(), EYsonType::MapFragment); + EXPECT_EQ(view.AsStringBuf(), "\"key\"=\"value\""); +} + +TEST(TStructuredPayloadTest, EmptyFragment) +{ + auto payload = MakeStructuredPayloadFromYson(MakeMapFragment("")); + EXPECT_EQ(GetYsonFromStructuredPayload(payload).AsStringBuf(), ""); +} + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace +} // namespace NYT::NLogging diff --git a/library/cpp/yt/logging/unittests/tagged_payload_ut.cpp b/library/cpp/yt/logging/unittests/tagged_payload_ut.cpp new file mode 100644 index 00000000000..57d9ac168f4 --- /dev/null +++ b/library/cpp/yt/logging/unittests/tagged_payload_ut.cpp @@ -0,0 +1,146 @@ +#include "helpers.h" + +#include <library/cpp/testing/gtest/gtest.h> + +#include <library/cpp/yt/logging/tagged_payload.h> + +#include <string> +#include <vector> + +namespace NYT::NLogging { +namespace { + +using namespace NDetail; + +//////////////////////////////////////////////////////////////////////////////// + +struct TDecodedPayload +{ + std::string Message; + std::vector<std::pair<std::string, std::string>> Tags; +}; + +TTaggedLogEventPayload Encode(TStringBuf message, const std::vector<std::pair<TStringBuf, TStringBuf>>& tags = {}) +{ + TTaggedPayloadWriter writer; + WriteMessage(&writer, message); + for (auto [key, value] : tags) { + WriteTag(&writer, key, value); + } + return writer.Finish(); +} + +TDecodedPayload Decode(const TTaggedLogEventPayload& payload) +{ + TTaggedPayloadReader reader(payload); + TDecodedPayload result; + result.Message = reader.ReadMessage(); + while (auto tag = reader.TryReadTag()) { + result.Tags.emplace_back(tag->Key, tag->Value); + } + return result; +} + +TEST(TTaggedPayloadTest, MessageOnly) +{ + auto decoded = Decode(Encode("Hello")); + EXPECT_EQ(decoded.Message, "Hello"); + EXPECT_TRUE(decoded.Tags.empty()); +} + +TEST(TTaggedPayloadTest, EmptyMessage) +{ + auto decoded = Decode(Encode("")); + EXPECT_EQ(decoded.Message, ""); + EXPECT_TRUE(decoded.Tags.empty()); +} + +TEST(TTaggedPayloadTest, OneTag) +{ + auto decoded = Decode(Encode("Message", {{"Key", "Value"}})); + EXPECT_EQ(decoded.Message, "Message"); + ASSERT_EQ(decoded.Tags.size(), 1u); + EXPECT_EQ(decoded.Tags[0].first, "Key"); + EXPECT_EQ(decoded.Tags[0].second, "Value"); +} + +TEST(TTaggedPayloadTest, ManyTags) +{ + auto decoded = Decode(Encode("Message", {{"Arg1", "123"}, {"Arg2", "test"}, {"Arg3", ""}})); + EXPECT_EQ(decoded.Message, "Message"); + ASSERT_EQ(decoded.Tags.size(), 3u); + EXPECT_EQ(decoded.Tags[0], std::pair(std::string("Arg1"), std::string("123"))); + EXPECT_EQ(decoded.Tags[1], std::pair(std::string("Arg2"), std::string("test"))); + EXPECT_EQ(decoded.Tags[2], std::pair(std::string("Arg3"), std::string(""))); +} + +TEST(TTaggedPayloadTest, BinarySafeValues) +{ + // Keys/values may carry arbitrary bytes, including embedded NULs and delimiters. + std::string message("a\0b ()", 6); + std::string key("k\0:", 3); + std::string value("v\0, ", 4); + + auto decoded = Decode(Encode(message, {{key, value}})); + EXPECT_EQ(decoded.Message, message); + ASSERT_EQ(decoded.Tags.size(), 1u); + EXPECT_EQ(decoded.Tags[0].first, key); + EXPECT_EQ(decoded.Tags[0].second, value); +} + +TEST(TTaggedPayloadTest, ReaderViewsPointIntoPayload) +{ + auto payload = Encode("Message", {{"Key", "Value"}}); + + auto pointsInto = [&] (TStringBuf view) { + return view.data() >= payload.Underlying().Begin() && view.data() + view.size() <= payload.Underlying().End(); + }; + + TTaggedPayloadReader reader(payload); + EXPECT_TRUE(pointsInto(reader.ReadMessage())); + auto tag = reader.TryReadTag(); + ASSERT_TRUE(tag.has_value()); + EXPECT_TRUE(pointsInto(tag->Key)); + EXPECT_TRUE(pointsInto(tag->Value)); +} + +TEST(TTaggedPayloadTest, WellKnownTag) +{ + TTaggedPayloadWriter writer; + WriteMessage(&writer, "Message"); + WriteTag(&writer, "Key", "Value"); + WriteWellKnownTag(&writer, "Error", "boom"); + auto payload = writer.Finish(); + + TTaggedPayloadReader reader(payload); + EXPECT_EQ(reader.ReadMessage(), "Message"); + + auto regular = reader.TryReadTag(); + ASSERT_TRUE(regular); + EXPECT_EQ(regular->Key, "Key"); + EXPECT_EQ(regular->Value, "Value"); + EXPECT_FALSE(regular->IsWellKnown); + + auto wellKnown = reader.TryReadTag(); + ASSERT_TRUE(wellKnown); + EXPECT_EQ(wellKnown->Key, "Error"); + EXPECT_EQ(wellKnown->Value, "boom"); + EXPECT_TRUE(wellKnown->IsWellKnown); + + EXPECT_FALSE(reader.TryReadTag()); +} + +TEST(TTaggedPayloadTest, FormatWellKnownTagTrailing) +{ + TTaggedPayloadWriter writer; + WriteMessage(&writer, "Message"); + WriteTag(&writer, "Key", "Value"); + WriteWellKnownTag(&writer, "Error", "boom"); + // Regular tags stay inline; the well-known tag is appended after the |(...)| group. + EXPECT_EQ(FormatTaggedPayload(writer.Finish()), "Message (Key: Value)\nboom"); +} + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace +} // namespace NYT::NLogging diff --git a/library/cpp/yt/logging/unittests/ya.make b/library/cpp/yt/logging/unittests/ya.make index 01946699bb2..cc068fa24c6 100644 --- a/library/cpp/yt/logging/unittests/ya.make +++ b/library/cpp/yt/logging/unittests/ya.make @@ -3,7 +3,10 @@ GTEST(unittester-library-logging) INCLUDE(${ARCADIA_ROOT}/library/cpp/yt/ya_cpp.make.inc) SRCS( + helpers.cpp logger_ut.cpp + tagged_payload_ut.cpp + structured_payload_ut.cpp static_analysis_ut.cpp ) diff --git a/library/cpp/yt/logging/ya.make b/library/cpp/yt/logging/ya.make index b9cf5c6e715..d839a6f07b7 100644 --- a/library/cpp/yt/logging/ya.make +++ b/library/cpp/yt/logging/ya.make @@ -4,6 +4,8 @@ INCLUDE(${ARCADIA_ROOT}/library/cpp/yt/ya_cpp.make.inc) SRCS( logger.cpp + tagged_payload.cpp + structured_payload.cpp ) PEERDIR( @@ -46,33 +46,33 @@ REGISTRY_ENDPOINT = os.environ.get("YA_REGISTRY_ENDPOINT", "https://devtools-reg PLATFORM_MAP = { "data": { "win32": { - "md5": "0080c0947dc0ddf9b20ca50821b32356", + "md5": "9dee86473a0e4b0c0a0d963f419dbf33", "urls": [ - f"{REGISTRY_ENDPOINT}/12778691402" + f"{REGISTRY_ENDPOINT}/12807581301" ] }, "darwin": { - "md5": "257864a1bc4ab2f0beef69bb19c5944f", + "md5": "e5067a1d83c2dd418be0faf32c2ab354", "urls": [ - f"{REGISTRY_ENDPOINT}/12778688709" + f"{REGISTRY_ENDPOINT}/12807578337" ] }, "darwin-arm64": { - "md5": "8166643048c1c44d3b172d57f9f558b4", + "md5": "732315791641f6d8dec06aad3ac70427", "urls": [ - f"{REGISTRY_ENDPOINT}/12778685977" + f"{REGISTRY_ENDPOINT}/12807575686" ] }, "linux-aarch64": { - "md5": "0cb28f3d593d8e7751db41688da47d48", + "md5": "9ac4edc8c328aff5bbb941668634ab70", "urls": [ - f"{REGISTRY_ENDPOINT}/12778683580" + f"{REGISTRY_ENDPOINT}/12807572366" ] }, "linux": { - "md5": "eada2433c6c25c0096f9c77433a82fdf", + "md5": "01654ceac32bcf11ad928ce31db2cc42", "urls": [ - f"{REGISTRY_ENDPOINT}/12778694423" + f"{REGISTRY_ENDPOINT}/12807584185" ] } } diff --git a/yql/essentials/ast/yql_constraint.cpp b/yql/essentials/ast/yql_constraint.cpp index 0946468b2a5..7fe88ae991a 100644 --- a/yql/essentials/ast/yql_constraint.cpp +++ b/yql/essentials/ast/yql_constraint.cpp @@ -2026,7 +2026,7 @@ NYT::TNode TStreamingConstraintNode::ToYson() const { } bool TStreamingConstraintNode::IsApplicableToType(const TTypeAnnotationNode& type) const { - return IsIn({ETypeAnnotationKind::List, ETypeAnnotationKind::Stream, ETypeAnnotationKind::Flow}, type.GetKind()); + return IsIn({ETypeAnnotationKind::Tuple, ETypeAnnotationKind::List, ETypeAnnotationKind::Stream, ETypeAnnotationKind::Flow}, type.GetKind()); } ////////////////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/yql/essentials/core/ut/yql_expr_constraint_ut.cpp b/yql/essentials/core/ut/yql_expr_constraint_ut.cpp index 3b23e8d8385..cbbcf49d481 100644 --- a/yql/essentials/core/ut/yql_expr_constraint_ut.cpp +++ b/yql/essentials/core/ut/yql_expr_constraint_ut.cpp @@ -4199,6 +4199,79 @@ Y_UNIT_TEST(StreamingConstraintShuffleByKeys) { } } +Y_UNIT_TEST(DirectStreamingConstraintWithSwitch) { + const auto s = R"(( +(let res (DataSink 'result)) + +(let list1 (AsList + (AsStruct '('key (String '4)) '('subkey (String 'c)) '('value (String 'v))) + (AsStruct '('key (String '1)) '('subkey (String 'd)) '('value (String 'v))) + (AsStruct '('key (String '3)) '('subkey (String 'b)) '('value (String 'v))) +)) + +(let list2 (AsList + (AsStruct '('key (String '2)) '('subkey (String 'c)) '('value (String 'v))) + (AsStruct '('key (String '5)) '('subkey (String 'd)) '('value (String 'v))) + (AsStruct '('key (String '4)) '('subkey (String 'b)) '('value (String 'v))) +)) + +(let data (Mux '(list1 list2))) +(let data (AssumeConstraints data '"{\"Streaming\" = #}")) + +(let data (Switch (Iterator data) '0 '('0) (lambda '(s) (FlatMap s (lambda '(item) (Just item)))) '('1) (lambda '(s) s))) + +(let result (Nth (Demux data) '1)) + +(let world (Write! world res (Key) (Collect result) '())) +(let world (Commit! world res)) +(return world) +))"; + + TExprContext exprCtx; + const auto exprRoot = ParseAndAnnotate(s, exprCtx); + CheckConstraint<TStreamingConstraintNode>(exprRoot, "Switch", "Streaming"); + CheckConstraint<TStreamingConstraintNode>(exprRoot, "FlatMap", "Streaming"); + CheckConstraint<TStreamingConstraintNode>(exprRoot, "Demux", "Streaming"); + CheckConstraint<TStreamingConstraintNode>(exprRoot, "Nth", "Streaming"); +} + +Y_UNIT_TEST(MultiItemStreamingConstraintWithSwitch) { + const auto s = R"(( +(let res (DataSink 'result)) + +(let list1 (AsList + (AsStruct '('key (String '4)) '('subkey (String 'c)) '('value (String 'v))) + (AsStruct '('key (String '1)) '('subkey (String 'd)) '('value (String 'v))) + (AsStruct '('key (String '3)) '('subkey (String 'b)) '('value (String 'v))) +)) + +(let list2 (AsList + (AsStruct '('key (String '2)) '('subkey (String 'c)) '('value (String 'v))) + (AsStruct '('key (String '5)) '('subkey (String 'd)) '('value (String 'v))) + (AsStruct '('key (String '4)) '('subkey (String 'b)) '('value (String 'v))) +)) +(let list2 (AssumeConstraints list2 '"{\"Streaming\" = #}")) + +(let data (Mux '(list1 list2))) + +(let data (Switch (Iterator data) '0 '('0) (lambda '(s) (FlatMap s (lambda '(item) (Just item)))) '('1) (lambda '(s) s))) + +(let result (Nth (Demux data) '1)) + +(let world (Write! world res (Key) (Collect result) '())) +(let world (Commit! world res)) +(return world) +))"; + + TExprContext exprCtx; + const auto exprRoot = ParseAndAnnotate(s, exprCtx); + CheckConstraint<TStreamingConstraintNode>(exprRoot, "Mux", "Streaming"); + CheckConstraint<TStreamingConstraintNode>(exprRoot, "Switch", "Streaming"); + CheckConstraint<TStreamingConstraintNode>(exprRoot, "FlatMap", "Streaming"); + CheckConstraint<TStreamingConstraintNode>(exprRoot, "Demux", "Streaming"); + CheckConstraint<TStreamingConstraintNode>(exprRoot, "Nth", "Streaming"); +} + } // Y_UNIT_TEST_SUITE(TYqlExprConstraints) } // namespace NYql diff --git a/yql/essentials/core/yql_expr_constraint.cpp b/yql/essentials/core/yql_expr_constraint.cpp index 864f46e248d..eb8a9694f63 100644 --- a/yql/essentials/core/yql_expr_constraint.cpp +++ b/yql/essentials/core/yql_expr_constraint.cpp @@ -207,6 +207,7 @@ public: Functions_["DynamicVariant"] = &TCallableConstraintTransformer::DynamicVariantWrap; Functions_["Guess"] = &TCallableConstraintTransformer::GuessWrap; Functions_["Mux"] = &TCallableConstraintTransformer::MuxWrap; + Functions_["Demux"] = &TCallableConstraintTransformer::FromFirst<TStreamingConstraintNode>; Functions_["Nth"] = &TCallableConstraintTransformer::NthWrap; Functions_["EquiJoin"] = &TCallableConstraintTransformer::EquiJoinWrap; Functions_["JoinDict"] = &TCallableConstraintTransformer::JoinDictWrap; @@ -1844,6 +1845,8 @@ private: TStatus SwitchWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const { TStatus status = TStatus::Ok; TDynBitMap outFromChildren; // children, from which take a multi constraint for output + const auto* inputStreaming = input->Head().GetConstraint<TStreamingConstraintNode>(); + if (const auto multi = input->Head().GetConstraint<TMultiConstraintNode>()) { for (size_t i = 2; i < input->ChildrenSize(); ++i) { TMultiConstraintNode::TMapType items; @@ -1861,18 +1864,33 @@ private: if (input->Child(i)->ChildrenSize() > 1) { argConstraints.push_back(ctx.MakeConstraint<TMultiConstraintNode>(std::move(items))); argConstraints.push_back(ctx.MakeConstraint<TVarIndexConstraintNode>(input->Child(i)->ChildrenSize())); + + if (inputStreaming) { + argConstraints.push_back(inputStreaming); + } } else { argConstraints = items.front().second.GetAllConstraints(); + + if (inputStreaming && !items.front().second.GetConstraint<TStreamingConstraintNode>()) { + argConstraints.push_back(inputStreaming); + } } + } else if (inputStreaming) { + argConstraints.push_back(inputStreaming); } status = status.Combine(UpdateLambdaConstraints(input->ChildRef(++i), ctx, {argConstraints})); } } else { - const bool inVar = GetSeqItemType(*input->Head().GetTypeAnn()).GetKind() == ETypeAnnotationKind::Variant; - const TSmallVec<TConstraintNode::TListType> argConstraints(1U, inVar ? TConstraintNode::TListType() : input->Head().GetAllConstraints()); + TConstraintNode::TListType argConstraintsList; + if (GetSeqItemType(*input->Head().GetTypeAnn()).GetKind() != ETypeAnnotationKind::Variant) { + argConstraintsList = input->Head().GetAllConstraints(); + } else if (inputStreaming) { + argConstraintsList.push_back(inputStreaming); + } + for (size_t i = 3; i < input->ChildrenSize(); i += 2) { - status = status.Combine(UpdateLambdaConstraints(input->ChildRef(i), ctx, argConstraints)); + status = status.Combine(UpdateLambdaConstraints(input->ChildRef(i), ctx, {argConstraintsList})); } outFromChildren.Set(0, input->ChildrenSize()); } @@ -1883,6 +1901,7 @@ private: const auto inputVarIndex = input->Head().GetConstraint<TVarIndexConstraintNode>(); const bool emptyInput = input->Head().GetConstraint<TEmptyConstraintNode>(); + bool hasStreamingOutput = false; if (GetSeqItemType(*input->GetTypeAnn()).GetKind() == ETypeAnnotationKind::Variant) { ui32 outIndexOffset = 0; TMultiConstraintNode::TMapType multiItems; @@ -1948,6 +1967,9 @@ private: } if (!multiItems.empty()) { + hasStreamingOutput = std::any_of(multiItems.begin(), multiItems.end(), [](const auto& item) { + return !!item.second.template GetConstraint<TStreamingConstraintNode>(); + }); input->AddConstraint(ctx.MakeConstraint<TMultiConstraintNode>(std::move(multiItems))); } if (emptyOut) { @@ -1957,6 +1979,11 @@ private: YQL_ENSURE(input->ChildrenSize() == 4); input->CopyConstraints(*input->Child(3)); } + + if (hasStreamingOutput || inputStreaming) { + input->AddConstraint(ctx.MakeConstraint<TStreamingConstraintNode>()); + } + return FromFirst<TEmptyConstraintNode>(input, output, ctx); } @@ -2234,19 +2261,26 @@ private: TMultiConstraintNode::TMapType items; ui32 index = 0; ui32 emptyCount = 0; + bool hasStreaming = false; for (auto& child: input->Head().Children()) { if (!child->GetConstraint<TEmptyConstraintNode>()) { items.push_back(std::make_pair(index, child->GetConstraintSet())); + hasStreaming = hasStreaming || !!child->GetConstraint<TStreamingConstraintNode>(); } else { ++emptyCount; } ++index; } + if (!items.empty()) { input->AddConstraint(ctx.MakeConstraint<TMultiConstraintNode>(std::move(items))); } else if (index == emptyCount) { input->AddConstraint(ctx.MakeConstraint<TEmptyConstraintNode>()); } + + if (hasStreaming) { + input->AddConstraint(ctx.MakeConstraint<TStreamingConstraintNode>()); + } } } return TStatus::Ok; @@ -2289,6 +2323,10 @@ private: input->SetConstraints(*c); } } + + if (const auto* c = input->Head().GetConstraint<TStreamingConstraintNode>()) { + input->AddConstraint(c); + } } return TStatus::Ok; } diff --git a/yql/essentials/public/purecalc/io_specs/protobuf/no_llvm/ya.make b/yql/essentials/public/purecalc/io_specs/protobuf/no_llvm/ya.make new file mode 100644 index 00000000000..76ed1dc0196 --- /dev/null +++ b/yql/essentials/public/purecalc/io_specs/protobuf/no_llvm/ya.make @@ -0,0 +1,23 @@ +LIBRARY() + +ADDINCL( + yql/essentials/public/purecalc/io_specs/protobuf +) + +SRCDIR( + yql/essentials/public/purecalc/io_specs/protobuf +) + +SRCS( + spec.cpp + proto_variant.cpp +) + +PEERDIR( + yql/essentials/public/purecalc/common/no_llvm + yql/essentials/public/purecalc/io_specs/protobuf_raw/no_llvm +) + +YQL_LAST_ABI_VERSION() + +END() diff --git a/yql/essentials/public/purecalc/io_specs/protobuf/ya.make b/yql/essentials/public/purecalc/io_specs/protobuf/ya.make index b9441ceecf4..a7cca185da2 100644 --- a/yql/essentials/public/purecalc/io_specs/protobuf/ya.make +++ b/yql/essentials/public/purecalc/io_specs/protobuf/ya.make @@ -14,6 +14,10 @@ YQL_LAST_ABI_VERSION() END() +RECURSE( + no_llvm +) + RECURSE_FOR_TESTS( ut ) diff --git a/yql/essentials/public/purecalc/io_specs/protobuf_raw/no_llvm/ya.make b/yql/essentials/public/purecalc/io_specs/protobuf_raw/no_llvm/ya.make new file mode 100644 index 00000000000..f89c7800864 --- /dev/null +++ b/yql/essentials/public/purecalc/io_specs/protobuf_raw/no_llvm/ya.make @@ -0,0 +1,23 @@ +LIBRARY() + +ADDINCL( + yql/essentials/public/purecalc/io_specs/protobuf_raw +) + +SRCDIR( + yql/essentials/public/purecalc/io_specs/protobuf_raw +) + +SRCS( + proto_holder.cpp + spec.cpp +) + +PEERDIR( + yql/essentials/public/purecalc/common/no_llvm + yql/essentials/public/purecalc/helpers/protobuf +) + +YQL_LAST_ABI_VERSION() + +END() diff --git a/yql/essentials/public/purecalc/io_specs/protobuf_raw/ya.make b/yql/essentials/public/purecalc/io_specs/protobuf_raw/ya.make index db3fab7e7a5..4faaf4d2de1 100644 --- a/yql/essentials/public/purecalc/io_specs/protobuf_raw/ya.make +++ b/yql/essentials/public/purecalc/io_specs/protobuf_raw/ya.make @@ -14,3 +14,7 @@ SRCS( YQL_LAST_ABI_VERSION() END() + +RECURSE( + no_llvm +) diff --git a/yql/essentials/sql/v1/SQLv1Antlr4.g.in b/yql/essentials/sql/v1/SQLv1Antlr4.g.in index d1eba59441a..54ffbdeec54 100644 --- a/yql/essentials/sql/v1/SQLv1Antlr4.g.in +++ b/yql/essentials/sql/v1/SQLv1Antlr4.g.in @@ -882,6 +882,7 @@ column_option: | default_value | compression | encoding + | generated_always ; compression: COMPRESSION LPAREN (compression_setting_entry (COMMA compression_setting_entry)*)? COMMA? RPAREN; @@ -896,6 +897,7 @@ encoding_config: encoding_config_name (LPAREN (encoding_setting_entry (COMMA enc encoding_config_name: an_id_or_type; encoding_setting_entry: an_id EQUALS encoding_setting_value; encoding_setting_value: integer | id; +generated_always: (GENERATED ALWAYS)? AS LPAREN expr RPAREN (STORED | VIRTUAL)?; column_order_by_specification: an_id (ASC | DESC)?; @@ -1476,6 +1478,7 @@ keyword_as_compat: | ADD | AFTER | ALTER + | ALWAYS | ANALYZE | AND | ANSI @@ -1551,6 +1554,7 @@ keyword_as_compat: | FOR | FOREIGN | FUNCTION + | GENERATED | GLOB | GLOBAL | GRANT @@ -1669,6 +1673,7 @@ keyword_as_compat: | SOURCE | START | STATISTICS + | STORED | STREAMING | SUBQUERY | SUBSET @@ -1717,6 +1722,7 @@ keyword_compat: ( | ADD | AFTER | ALTER + | ALWAYS | ANALYZE | AND | ANSI @@ -1792,6 +1798,7 @@ keyword_compat: ( | FOR | FOREIGN | FUNCTION + | GENERATED | GLOB | GLOBAL | GRANT @@ -1910,6 +1917,7 @@ keyword_compat: ( | SOURCE | START | STATISTICS + | STORED | STREAMING | SUBQUERY | SUBSET @@ -2055,6 +2063,7 @@ ADD: A D D; AFTER: A F T E R; ALL: A L L; ALTER: A L T E R; +ALWAYS: A L W A Y S; ANALYZE: A N A L Y Z E; AND: A N D; ANSI: A N S I; @@ -2159,6 +2168,7 @@ FOREIGN: F O R E I G N; FROM: F R O M; FULL: F U L L; FUNCTION: F U N C T I O N; +GENERATED: G E N E R A T E D; GLOB: G L O B; GLOBAL: G L O B A L; GRANT: G R A N T; @@ -2300,6 +2310,7 @@ SECRET: S E C R E T; SEQUENCE: S E Q U E N C E; SOURCE: S O U R C E; START: S T A R T; +STORED: S T O R E D; STATISTICS: S T A T I S T I C S; STREAM: S T R E A M; STREAMING: S T R E A M I N G; diff --git a/yql/essentials/sql/v1/format/sql_format_ut.h b/yql/essentials/sql/v1/format/sql_format_ut.h index ce5d448c9dd..6406608c26b 100644 --- a/yql/essentials/sql/v1/format/sql_format_ut.h +++ b/yql/essentials/sql/v1/format/sql_format_ut.h @@ -432,6 +432,20 @@ Y_UNIT_TEST(CreateTable) { {"create table\tuser(key int32, val String encoding(off))", "CREATE TABLE user (\n\tkey int32,\n\tval String ENCODING (off)\n);\n"}, {"create table\tuser(key int32, val String encoding())", "CREATE TABLE user (\n\tkey int32,\n\tval String ENCODING ()\n);\n"}, {"create table user(key int32, val String encoding(dict(max_size=100)))", "CREATE TABLE user (\n\tkey int32,\n\tval String ENCODING (dict (max_size = 100))\n);\n"}, + {"create table user(key int32, val int64 generated always as (key+1) stored)", + "CREATE TABLE user (\n\tkey int32,\n\tval int64 GENERATED ALWAYS AS (key + 1) STORED\n);\n"}, + {"create table user(key int32, val int64 generated always as (key+1) virtual)", + "CREATE TABLE user (\n\tkey int32,\n\tval int64 GENERATED ALWAYS AS (key + 1) VIRTUAL\n);\n"}, + {"create table user(key int32, val int64 generated always as (key+1))", + "CREATE TABLE user (\n\tkey int32,\n\tval int64 GENERATED ALWAYS AS (key + 1)\n);\n"}, + {"create table user(key int32, val int64 as (key+1) stored)", + "CREATE TABLE user (\n\tkey int32,\n\tval int64 AS (key + 1) STORED\n);\n"}, + {"create table user(key int32, val int64 as (key+1))", + "CREATE TABLE user (\n\tkey int32,\n\tval int64 AS (key + 1)\n);\n"}, + {"create table user(key int32, val int64 GeNeRaTeD AlWaYs As (key+1) StOrEd)", + "CREATE TABLE user (\n\tkey int32,\n\tval int64 GENERATED ALWAYS AS (key + 1) STORED\n);\n"}, + {"create table user(key int32, val int64 (not null, generated always as (key+1) stored))", + "CREATE TABLE user (\n\tkey int32,\n\tval int64 (NOT NULL, GENERATED ALWAYS AS (key + 1) STORED)\n);\n"}, }; TSetup setup; @@ -691,6 +705,18 @@ Y_UNIT_TEST(AlterTable) { "ALTER TABLE t\n\tALTER COLUMN c SET ENCODING ()\n;\n"}, {"alter table t alter column c set encoding(dict(max_size=100))", "ALTER TABLE t\n\tALTER COLUMN c SET ENCODING (dict (max_size = 100))\n;\n"}, + {"alter table user add column val int64 generated always as (key+1) stored", + "ALTER TABLE user\n\tADD COLUMN val int64 GENERATED ALWAYS AS (key + 1) STORED\n;\n"}, + {"alter table user add column val int64 generated always as (key+1) virtual", + "ALTER TABLE user\n\tADD COLUMN val int64 GENERATED ALWAYS AS (key + 1) VIRTUAL\n;\n"}, + {"alter table user add column val int64 generated always as (key+1)", + "ALTER TABLE user\n\tADD COLUMN val int64 GENERATED ALWAYS AS (key + 1)\n;\n"}, + {"alter table user add column val int64 as (key+1) stored", + "ALTER TABLE user\n\tADD COLUMN val int64 AS (key + 1) STORED\n;\n"}, + {"alter table user add val int64 GeNeRaTeD AlWaYs As (key+1) StOrEd", + "ALTER TABLE user\n\tADD val int64 GENERATED ALWAYS AS (key + 1) STORED\n;\n"}, + {"alter table user add column val int64 (not null, generated always as (key+1) stored)", + "ALTER TABLE user\n\tADD COLUMN val int64 (NOT NULL, GENERATED ALWAYS AS (key + 1) STORED)\n;\n"}, }; TSetup setup; diff --git a/yql/essentials/sql/v1/node.h b/yql/essentials/sql/v1/node.h index c9e8b16a799..b79eb7d04bc 100644 --- a/yql/essentials/sql/v1/node.h +++ b/yql/essentials/sql/v1/node.h @@ -791,12 +791,19 @@ struct TEncoding { TMap<TString, TNodePtr> Entries; }; +struct TGeneratedColumn { + TString ContextPrefix; + TString ExprBody; + bool Stored = false; +}; + struct TColumnOptions { TNodePtr DefaultExpr; TVector<TIdentifier> Families; TMaybe<TCompression> Compression; bool Nullable = true; TMaybe<TVector<TEncoding>> ColumnEncoding; + TMaybe<TGeneratedColumn> Generated; }; struct TColumnSchema { @@ -821,6 +828,7 @@ struct TColumnSchema { bool Nullable = false; bool Serial = false; TMaybe<TVector<TEncoding>> ColumnEncoding; + TMaybe<TGeneratedColumn> Generated; }; struct TColumns: public TSimpleRefCount<TColumns> { diff --git a/yql/essentials/sql/v1/query.cpp b/yql/essentials/sql/v1/query.cpp index 6ec40c25ce5..0276df12534 100644 --- a/yql/essentials/sql/v1/query.cpp +++ b/yql/essentials/sql/v1/query.cpp @@ -490,6 +490,13 @@ INode::TPtr CreateEncodingsListSettings(const TVector<TEncoding>& columnEncoding return encodingsList; } +INode::TPtr CreateGeneratedValue(const TGeneratedColumn& generated, TPosition pos, const INode& node) { + const ui32 flags = NYql::TAstNodeFlags::ArbitraryContent | NYql::TAstNodeFlags::UnstableFormat; + auto contextText = BuildQuotedAtom(pos, generated.ContextPrefix, flags); + auto exprText = BuildQuotedAtom(pos, generated.ExprBody, flags); + return node.Q(node.Y(contextText, exprText, node.Q(generated.Stored ? "stored" : "virtual"))); +} + } // namespace class TPrepTableKeys: public ITableKeys { @@ -1357,6 +1364,11 @@ public: columnConstraints = L(columnConstraints, Q(Y(Q("default"), col.DefaultExpr))); } + if (col.Generated) { + auto generatedValue = CreateGeneratedValue(*col.Generated, Pos_, *this); + columnConstraints = L(columnConstraints, Q(Y(Q("generated"), generatedValue))); + } + columnDesc = L(columnDesc, Q(Y(Q("columnConstrains"), Q(columnConstraints)))); if (col.Compression) { @@ -1715,6 +1727,11 @@ public: columnConstraints = L(columnConstraints, Q(Y(Q("default"), col.DefaultExpr))); } + if (col.Generated) { + auto generatedValue = CreateGeneratedValue(*col.Generated, Pos_, *this); + columnConstraints = L(columnConstraints, Q(Y(Q("generated"), generatedValue))); + } + columnDesc = L(columnDesc, Q(Y(Q("columnConstrains"), Q(columnConstraints)))); if (col.Compression) { diff --git a/yql/essentials/sql/v1/sql_query.cpp b/yql/essentials/sql/v1/sql_query.cpp index c6de747288b..42228fa5586 100644 --- a/yql/essentials/sql/v1/sql_query.cpp +++ b/yql/essentials/sql/v1/sql_query.cpp @@ -2778,6 +2778,17 @@ bool TSqlQuery::AlterTableAddColumn(const TRule_alter_table_add_column& node, TA Ctx_.Error() << "Several column families for a single column are not yet supported"; return false; } + if (columnSchema->Generated) { + const TString& service = Ctx_.Scoped->CurrService; + if (service != KikimrProviderName && service != YdbProviderName) { + Ctx_.Error(columnSchema->Pos) << "GENERATED ALWAYS AS columns are supported only for the ydb provider"; + return false; + } + if (params.TableType != ETableType::Table) { + Ctx_.Error(columnSchema->Pos) << "GENERATED ALWAYS AS columns are supported only for CREATE TABLE and ALTER TABLE"; + return false; + } + } params.AddColumns.push_back(*columnSchema); return true; } @@ -5020,6 +5031,11 @@ bool BuildColumnFeatures(std::map<TString, TDeferredAtom>& result, const TRule_c return false; } + if (options->Generated) { + translation.Context().Error(pos) << "GENERATED ALWAYS AS columns are supported only for CREATE TABLE and ALTER TABLE"; + return false; + } + auto& typeBind = columnSchema.GetRule_type_name_or_bind2(); switch (typeBind.Alt_case()) { case TRule_type_name_or_bind::kAltTypeNameOrBind1: { diff --git a/yql/essentials/sql/v1/sql_translation.cpp b/yql/essentials/sql/v1/sql_translation.cpp index 141d416f6d4..c40afa7705e 100644 --- a/yql/essentials/sql/v1/sql_translation.cpp +++ b/yql/essentials/sql/v1/sql_translation.cpp @@ -17,6 +17,7 @@ #include <util/generic/scope.h> #include <util/string/join.h> +#include <util/string/strip.h> #include <library/cpp/protobuf/util/simple_reflection.h> @@ -1731,6 +1732,7 @@ TMaybe<TColumnOptions> ColumnOptions(const TRule_column_schema& node, TSqlTransl TMaybe<TCompression> compression; bool nullable = true; TMaybe<TVector<TEncoding>> columnEncoding; + TMaybe<TGeneratedColumn> generated; const auto& optionsList = node.GetRule_column_option_list3(); @@ -1740,9 +1742,10 @@ TMaybe<TColumnOptions> ColumnOptions(const TRule_column_schema& node, TSqlTransl DefaultValue, Compression, Encoding, + Generated, }; - TVector<TRule_column_option> columnOptions(Reserve(static_cast<size_t>(EOption::Encoding) + 1)); + TVector<TRule_column_option> columnOptions(Reserve(static_cast<size_t>(EOption::Generated) + 1)); { switch (optionsList.Alt_case()) { @@ -1873,6 +1876,57 @@ TMaybe<TColumnOptions> ColumnOptions(const TRule_column_schema& node, TSqlTransl break; } + case TRule_column_option::kAltColumnOption6: { // generated_always + const auto& opt = rule.GetAlt_column_option6().GetRule_generated_always1(); + if (std::find(usedOptions.begin(), usedOptions.end(), EOption::Generated) != usedOptions.end()) { + TPosition pos = ctx.Context().TokenPosition(opt.GetToken2()); + ctx.Context().Error(pos) << "'GENERATED ALWAYS AS' option can be specified only once"; + return {}; + } + + usedOptions.push_back(EOption::Generated); + + { + TColumnRefScope scope(ctx.Context(), EColumnRefState::Allow); + TSqlExpression expr(ctx); + if (!Unwrap(expr.Build(opt.GetRule_expr4()))) { + return {}; + } + } + + TStringBuilder contextPrefix; + if (!BuildContextRecreationQuery(ctx.Context(), contextPrefix)) { + return {}; + } + + const auto& lparen = opt.GetToken3(); + const auto& rparen = opt.GetToken5(); + auto begin = GetQueryPosition(ctx.Context().Query, lparen); + auto end = GetQueryPosition(ctx.Context().Query, rparen); + + if (begin == std::string::npos || end == std::string::npos) { + TPosition pos = ctx.Context().TokenPosition(opt.GetToken2()); + ctx.Context().Error(pos) << "Failed to extract GENERATED ALWAYS AS expression text"; + return {}; + } + + begin += lparen.value().size(); + TString exprBody = StripString(ctx.Context().Query.substr(begin, end - begin)); + + bool generatedStored = false; + if (opt.HasBlock6()) { + const auto tokenId = opt.GetBlock6().GetToken1().GetId(); + generatedStored = IS_TOKEN(tokenId, STORED); + } + + generated = TGeneratedColumn{ + .ContextPrefix = std::move(contextPrefix), + .ExprBody = std::move(exprBody), + .Stored = generatedStored, + }; + + break; + } case TRule_column_option::ALT_NOT_SET: YQL_ENSURE(false, "Unreachable"); } @@ -1885,6 +1939,7 @@ TMaybe<TColumnOptions> ColumnOptions(const TRule_column_schema& node, TSqlTransl .Compression = std::move(compression), .Nullable = nullable, .ColumnEncoding = std::move(columnEncoding), + .Generated = std::move(generated), }; } @@ -1899,7 +1954,7 @@ TMaybe<TColumnSchema> TSqlTranslation::ColumnSchemaImpl(const TRule_column_schem return {}; } - auto&& [defaultExpr, families, compression, nullable, columnEncoding] = columnOptions.GetRef(); + auto&& [defaultExpr, families, compression, nullable, columnEncoding, generated] = columnOptions.GetRef(); if (!type) { type = TypeNodeOrBind(node.GetRule_type_name_or_bind2()); @@ -1919,6 +1974,7 @@ TMaybe<TColumnSchema> TSqlTranslation::ColumnSchemaImpl(const TRule_column_schem .Nullable = nullable, .Serial = serial, .ColumnEncoding = columnEncoding, + .Generated = generated, }; } @@ -2056,6 +2112,17 @@ bool TSqlTranslation::CreateTableEntry(const TRule_create_table_entry& node, TCr Ctx_.Error() << "Several column families for a single column are not yet supported"; return false; } + if (columnSchema->Generated) { + const TString& service = Ctx_.Scoped->CurrService; + if (service != KikimrProviderName && service != YdbProviderName) { + Ctx_.Error(columnSchema->Pos) << "GENERATED ALWAYS AS columns are supported only for the ydb provider"; + return false; + } + if (params.TableType != ETableType::Table) { + Ctx_.Error(columnSchema->Pos) << "GENERATED ALWAYS AS columns are supported only for CREATE TABLE and ALTER TABLE"; + return false; + } + } params.Columns.push_back(*columnSchema); break; } diff --git a/yql/essentials/sql/v1/sql_ut_common.h b/yql/essentials/sql/v1/sql_ut_common.h index b53675b12ea..8558598e1da 100644 --- a/yql/essentials/sql/v1/sql_ut_common.h +++ b/yql/essentials/sql/v1/sql_ut_common.h @@ -16190,4 +16190,415 @@ Y_UNIT_TEST(OptionalNull) { } // Y_UNIT_TEST_SUITE(NullAsTypeName) +Y_UNIT_TEST_SUITE(GeneratedColumns) { + +Y_UNIT_TEST(CreateTableGeneratedVirtualByDefault) { + auto res = SqlToYql(R"sql( + USE ydb; + CREATE TABLE tbl ( + key Uint32, + gen Int32 AS (5), + PRIMARY KEY (key) + ); + )sql"); + UNIT_ASSERT_C(res.IsOk(), Err2Str(res)); + + const auto program = GetPrettyPrint(res); + UNIT_ASSERT_STRING_CONTAINS(program, R"#('('generated '('"USE ydb;\n" '"5" 'virtual)))#"); + UNIT_ASSERT(!program.Contains("stored")); + UNIT_ASSERT(!program.Contains(R"#((Int32 '"5"))#")); +} + +Y_UNIT_TEST(CreateTableGeneratedExplicitVirtual) { + auto res = SqlToYql(R"sql( + USE ydb; + CREATE TABLE tbl ( + key Uint32, + gen Int32 AS (5) VIRTUAL, + PRIMARY KEY (key) + ); + )sql"); + UNIT_ASSERT_C(res.IsOk(), Err2Str(res)); + + const auto program = GetPrettyPrint(res); + UNIT_ASSERT_STRING_CONTAINS(program, R"#('('generated '('"USE ydb;\n" '"5" 'virtual)))#"); + UNIT_ASSERT(!program.Contains("stored")); +} + +Y_UNIT_TEST(CreateTableGeneratedStored) { + auto res = SqlToYql(R"sql( + USE ydb; + CREATE TABLE tbl ( + key Uint32, + gen Int32 AS (5) STORED, + PRIMARY KEY (key) + ); + )sql"); + UNIT_ASSERT_C(res.IsOk(), Err2Str(res)); + + const auto program = GetPrettyPrint(res); + UNIT_ASSERT_STRING_CONTAINS(program, R"#('('generated '('"USE ydb;\n" '"5" 'stored)))#"); + UNIT_ASSERT(!program.Contains("virtual")); +} + +Y_UNIT_TEST(CreateTableGeneratedAlwaysKeyword) { + auto res = SqlToYql(R"sql( + USE ydb; + CREATE TABLE tbl ( + key Uint32, + gen Int32 GENERATED ALWAYS AS (5) STORED, + PRIMARY KEY (key) + ); + )sql"); + UNIT_ASSERT_C(res.IsOk(), Err2Str(res)); + + const auto program = GetPrettyPrint(res); + UNIT_ASSERT_STRING_CONTAINS(program, R"#('('generated '('"USE ydb;\n" '"5" 'stored)))#"); + UNIT_ASSERT(!program.Contains("virtual")); +} + +Y_UNIT_TEST(CreateTableGeneratedWithoutAlways) { + auto res = SqlToYql(R"sql( + USE ydb; + CREATE TABLE tbl ( + id Uint32, + Col2 Uint32, + Col3 Uint32, + Col1 Uint32 AS (2 + 3), + PRIMARY KEY (id) + ); + )sql"); + UNIT_ASSERT_C(res.IsOk(), Err2Str(res)); + + const auto program = GetPrettyPrint(res); + UNIT_ASSERT_STRING_CONTAINS(program, R"#('('generated '('"USE ydb;\n" '"2 + 3" 'virtual)))#"); + UNIT_ASSERT(!program.Contains("stored")); +} + +Y_UNIT_TEST(CreateTableGeneratedDeterministicExpr) { + auto res = SqlToYql(R"sql( + USE ydb; + CREATE TABLE tbl ( + key Uint32, + gen Int32 AS (1 + 2) VIRTUAL, + PRIMARY KEY (key) + ); + )sql"); + UNIT_ASSERT_C(res.IsOk(), Err2Str(res)); + + const auto program = GetPrettyPrint(res); + UNIT_ASSERT_STRING_CONTAINS(program, R"#('('generated '('"USE ydb;\n" '"1 + 2" 'virtual)))#"); + UNIT_ASSERT(!program.Contains("stored")); +} + +Y_UNIT_TEST(CreateTableGeneratedNonDeterministicExpr) { + auto res = SqlToYql(R"sql( + USE ydb; + CREATE TABLE tbl ( + key Uint32, + gen Uint64 AS (RandomNumber(1)) STORED, + PRIMARY KEY (key) + ); + )sql"); + UNIT_ASSERT_C(res.IsOk(), Err2Str(res)); + + const auto program = GetPrettyPrint(res); + UNIT_ASSERT_STRING_CONTAINS(program, R"#('('generated '('"USE ydb;\n" '"RandomNumber(1)" 'stored)))#"); + UNIT_ASSERT(!program.Contains("virtual")); +} + +Y_UNIT_TEST(CreateTableGeneratedReferencesColumns) { + auto res = SqlToYql(R"sql( + USE ydb; + CREATE TABLE tbl ( + a Int32, + b Int32, + gen Int32 GENERATED ALWAYS AS (a + b) STORED, + PRIMARY KEY (a) + ); + )sql"); + UNIT_ASSERT_C(res.IsOk(), Err2Str(res)); + + const auto program = GetPrettyPrint(res); + UNIT_ASSERT_STRING_CONTAINS(program, R"#('('generated '('"USE ydb;\n" '"a + b" 'stored)))#"); + UNIT_ASSERT(!program.Contains("virtual")); +} + +Y_UNIT_TEST(CreateTableMultipleGeneratedColumns) { + auto res = SqlToYql(R"sql( + USE ydb; + CREATE TABLE tbl ( + key Uint32, + x Int32 AS (1) VIRTUAL, + y Int32 AS (2) STORED, + PRIMARY KEY (key) + ); + )sql"); + UNIT_ASSERT_C(res.IsOk(), Err2Str(res)); + + const auto program = GetPrettyPrint(res); + UNIT_ASSERT_STRING_CONTAINS(program, R"#('('generated '('"USE ydb;\n" '"1" 'virtual)))#"); + UNIT_ASSERT_STRING_CONTAINS(program, R"#('('generated '('"USE ydb;\n" '"2" 'stored)))#"); +} + +Y_UNIT_TEST(AlterTableAddGeneratedVirtualColumn) { + auto res = SqlToYql(R"sql( + USE ydb; + ALTER TABLE tbl ADD COLUMN gen Int32 AS (5); + )sql"); + UNIT_ASSERT_C(res.IsOk(), Err2Str(res)); + + const auto program = GetPrettyPrint(res); + UNIT_ASSERT_STRING_CONTAINS(program, "addColumns"); + UNIT_ASSERT_STRING_CONTAINS(program, R"#('('generated '('"USE ydb;\n" '"5" 'virtual)))#"); + UNIT_ASSERT(!program.Contains("stored")); +} + +Y_UNIT_TEST(AlterTableAddGeneratedStoredColumn) { + auto res = SqlToYql(R"sql( + USE ydb; + ALTER TABLE tbl ADD COLUMN gen Int32 GENERATED ALWAYS AS (a + b) STORED; + )sql"); + UNIT_ASSERT_C(res.IsOk(), Err2Str(res)); + + const auto program = GetPrettyPrint(res); + UNIT_ASSERT_STRING_CONTAINS(program, "addColumns"); + UNIT_ASSERT_STRING_CONTAINS(program, R"#('('generated '('"USE ydb;\n" '"a + b" 'stored)))#"); + UNIT_ASSERT(!program.Contains("virtual")); + UNIT_ASSERT(!program.Contains("(Member row")); +} + +Y_UNIT_TEST(AlterTableAddGeneratedColumnWithNamedLambda) { + auto res = SqlToYql(R"sql( + USE ydb; + $f = ($x) -> ($x * 2); + ALTER TABLE tbl ADD COLUMN gen Int32 GENERATED ALWAYS AS ($f(a)) STORED; + )sql"); + UNIT_ASSERT_C(res.IsOk(), Err2Str(res)); + + const auto program = GetPrettyPrint(res); + UNIT_ASSERT_STRING_CONTAINS(program, "addColumns"); + UNIT_ASSERT_STRING_CONTAINS(program, R"#('('generated '('"USE ydb;\n$f = ($x) -> ($x * 2);\n" '"$f(a)" 'stored)))#"); + UNIT_ASSERT(!program.Contains("virtual")); +} + +Y_UNIT_TEST(CreateTableGeneratedColumnWithNamedNodeRef) { + auto res = SqlToYql(R"sql( + USE ydb; + $factor = 10; + CREATE TABLE tbl ( + key Uint32, + num Int32, + gen Int32 GENERATED ALWAYS AS (num * $factor) STORED, + PRIMARY KEY (key) + ); + )sql"); + UNIT_ASSERT_C(res.IsOk(), Err2Str(res)); + + const auto program = GetPrettyPrint(res); + UNIT_ASSERT_STRING_CONTAINS(program, R"#('('generated '('"USE ydb;\n$factor = 10;\n" '"num * $factor" 'stored)))#"); + UNIT_ASSERT(!program.Contains("virtual")); + UNIT_ASSERT(!program.Contains("Read")); +} + +Y_UNIT_TEST(CreateTableGeneratedColumnWithSubqueryNamedNodeRef) { + auto res = SqlToYql(R"sql( + USE ydb; + $ids = (SELECT id FROM other); + CREATE TABLE tbl ( + key Uint32, + num Int32, + gen Bool GENERATED ALWAYS AS (num IN $ids) STORED, + PRIMARY KEY (key) + ); + )sql"); + UNIT_ASSERT_C(res.IsOk(), Err2Str(res)); + + const auto program = GetPrettyPrint(res); + UNIT_ASSERT_STRING_CONTAINS(program, R"#('('generated '('"USE ydb;\n$ids = (SELECT id FROM other);\n" '"num IN $ids" 'stored)))#"); + UNIT_ASSERT(!program.Contains("virtual")); + UNIT_ASSERT(!program.Contains("Read")); +} + +Y_UNIT_TEST(CreateTableGeneratedColumnWithInlineSubquery) { + auto res = SqlToYql(R"sql( + USE ydb; + CREATE TABLE tbl ( + key Uint32, + num Int32, + gen Bool GENERATED ALWAYS AS (num IN (SELECT id FROM other)) STORED, + PRIMARY KEY (key) + ); + )sql"); + UNIT_ASSERT_C(res.IsOk(), Err2Str(res)); + + const auto program = GetPrettyPrint(res); + UNIT_ASSERT_STRING_CONTAINS(program, R"#('('generated '('"USE ydb;\n" '"num IN (SELECT id FROM other)" 'stored)))#"); + UNIT_ASSERT(!program.Contains("virtual")); + UNIT_ASSERT(!program.Contains("Read")); +} + +Y_UNIT_TEST(CreateTableGeneratedColumnWithScalarSubquery) { + auto res = SqlToYql(R"sql( + USE ydb; + CREATE TABLE tbl ( + key Uint32, + num Int32, + gen Int32 GENERATED ALWAYS AS ((SELECT MAX(id) FROM other)) STORED, + PRIMARY KEY (key) + ); + )sql"); + UNIT_ASSERT_C(res.IsOk(), Err2Str(res)); + + const auto program = GetPrettyPrint(res); + UNIT_ASSERT_STRING_CONTAINS(program, R"#('('generated '('"USE ydb;\n" '"(SELECT MAX(id) FROM other)" 'stored)))#"); + UNIT_ASSERT(!program.Contains("virtual")); + UNIT_ASSERT(!program.Contains("Read")); +} + +Y_UNIT_TEST(CreateTableGeneratedWithPragma) { + auto res = SqlToYql(R"sql( + USE ydb; + PRAGMA FlexibleTypes; + CREATE TABLE tbl ( + key Uint32, + num Int32, + gen Int32 GENERATED ALWAYS AS (num / 2) STORED, + PRIMARY KEY (key) + ); + )sql"); + UNIT_ASSERT_C(res.IsOk(), Err2Str(res)); + + const auto program = GetPrettyPrint(res); + UNIT_ASSERT_STRING_CONTAINS(program, R"#('('generated '('"USE ydb;\nPRAGMA FlexibleTypes;\n" '"num / 2" 'stored)))#"); + UNIT_ASSERT(!program.Contains("virtual")); +} + +Y_UNIT_TEST(CreateTableGeneratedNestedPath) { + auto res = SqlToYql(R"sql( + USE ydb; + CREATE TABLE `a/b/tbl` ( + key Uint32, + num Int32, + gen Int32 GENERATED ALWAYS AS (num / 2) STORED, + PRIMARY KEY (key) + ); + )sql"); + UNIT_ASSERT_C(res.IsOk(), Err2Str(res)); + + const auto program = GetPrettyPrint(res); + UNIT_ASSERT_STRING_CONTAINS(program, R"#('('generated '('"USE ydb;\n" '"num / 2" 'stored)))#"); +} + +Y_UNIT_TEST(AlterTableAddGeneratedColumnNestedPath) { + auto res = SqlToYql(R"sql( + USE ydb; + ALTER TABLE `a/b/tbl` ADD COLUMN gen Int32 GENERATED ALWAYS AS (a + b) STORED; + )sql"); + UNIT_ASSERT_C(res.IsOk(), Err2Str(res)); + + const auto program = GetPrettyPrint(res); + UNIT_ASSERT_STRING_CONTAINS(program, R"#('('generated '('"USE ydb;\n" '"a + b" 'stored)))#"); +} + +Y_UNIT_TEST(CreateTableStoreGeneratedColumnFails) { + auto res = SqlToYql(R"sql( + USE ydb; + CREATE TABLESTORE tbl ( + key Uint32 NOT NULL, + gen Int32 AS (5) STORED, + PRIMARY KEY (key) + ); + )sql"); + UNIT_ASSERT(!res.IsOk()); + UNIT_ASSERT_STRING_CONTAINS(Err2Str(res), "GENERATED ALWAYS AS columns are supported only for CREATE TABLE and ALTER TABLE"); +} + +Y_UNIT_TEST(CreateExternalTableGeneratedColumnFails) { + auto res = SqlToYql(R"sql( + USE ydb; + CREATE EXTERNAL TABLE tbl ( + a Int32, + gen Int32 AS (a) STORED + ) WITH ( + DATA_SOURCE="/Root/mydatasource", + LOCATION="/folder1/*" + ); + )sql"); + UNIT_ASSERT(!res.IsOk()); + UNIT_ASSERT_STRING_CONTAINS(Err2Str(res), "GENERATED ALWAYS AS columns are supported only for CREATE TABLE and ALTER TABLE"); +} + +Y_UNIT_TEST(AlterExternalTableAddGeneratedColumnFails) { + auto res = SqlToYql(R"sql( + USE ydb; + ALTER EXTERNAL TABLE tbl ADD COLUMN gen Int32 AS (5) STORED; + )sql"); + UNIT_ASSERT(!res.IsOk()); + UNIT_ASSERT_STRING_CONTAINS(Err2Str(res), "GENERATED ALWAYS AS columns are supported only for CREATE TABLE and ALTER TABLE"); +} + +Y_UNIT_TEST(AlterTableStoreAddGeneratedColumnFails) { + auto res = SqlToYql(R"sql( + USE ydb; + ALTER TABLESTORE tbl ADD COLUMN gen Int32 AS (5) STORED; + )sql"); + UNIT_ASSERT(!res.IsOk()); + UNIT_ASSERT_STRING_CONTAINS(Err2Str(res), "GENERATED ALWAYS AS columns are supported only for CREATE TABLE and ALTER TABLE"); +} + +Y_UNIT_TEST(CreateTableGeneratedColumnNonYdbProviderFails) { + auto res = SqlToYql(R"sql( + USE plato; + CREATE TABLE tbl ( + a Int32, + gen Int32 AS (a) STORED, + PRIMARY KEY (a) + ); + )sql"); + UNIT_ASSERT(!res.IsOk()); + UNIT_ASSERT_STRING_CONTAINS(Err2Str(res), "GENERATED ALWAYS AS columns are supported only for the ydb provider"); +} + +Y_UNIT_TEST(CreateTableGeneratedColumnNonYdbProviderVirtualFails) { + auto res = SqlToYql(R"sql( + USE plato; + CREATE TABLE tbl ( + key Uint32, + gen Int32 GENERATED ALWAYS AS (5), + PRIMARY KEY (key) + ); + )sql"); + UNIT_ASSERT(!res.IsOk()); + UNIT_ASSERT_STRING_CONTAINS(Err2Str(res), "GENERATED ALWAYS AS columns are supported only for the ydb provider"); +} + +Y_UNIT_TEST(AlterTableAddGeneratedColumnNonYdbProviderFails) { + auto res = SqlToYql(R"sql( + USE plato; + ALTER TABLE tbl ADD COLUMN gen Int32 AS (5) STORED; + )sql"); + UNIT_ASSERT(!res.IsOk()); + UNIT_ASSERT_STRING_CONTAINS(Err2Str(res), "ALTER TABLE is not supported for yt provider"); +} + +Y_UNIT_TEST(CreateTableGeneratedStringFunction) { + auto res = SqlToYql(R"sql( + USE ydb; + CREATE TABLE tbl ( + id Int32, + name String, + upper_name String GENERATED ALWAYS AS (String::AsciiToUpper(name)) STORED, + PRIMARY KEY (id) + ); + )sql"); + UNIT_ASSERT_C(res.IsOk(), Err2Str(res)); + + const auto program = GetPrettyPrint(res); + UNIT_ASSERT_STRING_CONTAINS(program, R"#('('generated '('"USE ydb;\n" '"String::AsciiToUpper(name)" 'stored)))#"); + UNIT_ASSERT(!program.Contains("virtual")); +} + +} // Y_UNIT_TEST_SUITE(GeneratedColumns) + // NOLINTEND(misc-definitions-in-headers) diff --git a/yt/cpp/mapreduce/interface/logging/yt_log.cpp b/yt/cpp/mapreduce/interface/logging/yt_log.cpp index 1db626adb04..9de369fdcb5 100644 --- a/yt/cpp/mapreduce/interface/logging/yt_log.cpp +++ b/yt/cpp/mapreduce/interface/logging/yt_log.cpp @@ -2,10 +2,14 @@ #include "logger.h" +#include <library/cpp/yt/logging/tagged_payload.h> + #include <util/generic/guid.h> #include <util/system/mutex.h> +#include <variant> + namespace NYT { using namespace NLogging; @@ -44,13 +48,14 @@ public: void Enqueue(TLogEvent&& event) override { if (auto logger = GetLogger()) { + auto message = GetMessageFromTaggedPayload(std::get<TTaggedLogEventPayload>(event.Payload)); LogMessage( logger, ToImplLevel(event.Level), ::TSourceLocation(event.SourceFile, event.SourceLine), "%.*s", - event.MessageRef.size(), - event.MessageRef.begin()); + message.size(), + message.begin()); } else if (auto* defaultLogManager = GetDefaultLogManager()) { defaultLogManager->Enqueue(std::move(event)); } diff --git a/yt/yql/providers/yt/codec/yt_codec_tz.h b/yt/yql/providers/yt/codec/yt_codec_tz.h index deadb9debb4..f6b23af95b2 100644 --- a/yt/yql/providers/yt/codec/yt_codec_tz.h +++ b/yt/yql/providers/yt/codec/yt_codec_tz.h @@ -25,7 +25,7 @@ void WriteTzNative(NCommon::TOutputBuf& buf, typename NUdf::TDataType<T>::TLayou buf.WriteMany((const char*)&tzId, sizeof(tzId)); } -// yt native table format (string with presorted timestamp concatenated with timezone name) +// yt native table format (string with presorted timestamp concatenated with presorted timezone id) // only tzId is validated template<typename T> bool ReadTzTableNative(NCommon::TInputBuf& buf, typename NUdf::TDataType<T>::TLayout& value, ui16& tzId) { @@ -33,7 +33,7 @@ bool ReadTzTableNative(NCommon::TInputBuf& buf, typename NUdf::TDataType<T>::TLa using TLayout = NUdf::TDataType<T>::TLayout; TStringBuf str = buf.ReadYtString(); - if (str.size() < sizeof(value)) { + if (str.size() != sizeof(value) + sizeof(tzId)) { return false; } @@ -46,11 +46,9 @@ bool ReadTzTableNative(NCommon::TInputBuf& buf, typename NUdf::TDataType<T>::TLa value = (TLayout)SwapBytes(data); str.Skip(sizeof(data)); - if (auto maybeTzId = NKikimr::NMiniKQL::FindTimezoneId(str)) { - tzId = *maybeTzId; - return true; - } - return false; + std::memcpy(&tzId, str.data(), sizeof(tzId)); + tzId = SwapBytes(tzId); + return NKikimr::NMiniKQL::IsValidTimezoneId(tzId); } template<typename T> @@ -63,13 +61,12 @@ void WriteTzTableNative(NCommon::TOutputBuf& buf, typename NUdf::TDataType<T>::T data = 0x80 ^ data; } - auto tzName = NKikimr::NMiniKQL::GetTimezoneIANAName(tzId); - - ui32 size = sizeof(data) + tzName.size();; + ui32 size = sizeof(data) + sizeof(tzId); buf.Write(NYson::NDetail::StringMarker); buf.WriteVarI32(size); buf.WriteMany((const char*)&data, sizeof(data)); - buf.WriteMany(tzName); + tzId = SwapBytes(tzId); + buf.WriteMany((const char*)&tzId, sizeof(tzId)); } namespace { diff --git a/yt/yql/providers/yt/provider/yql_yt_ytflow_integration.cpp b/yt/yql/providers/yt/provider/yql_yt_ytflow_integration.cpp index 9e4c87017fe..c733740bf06 100644 --- a/yt/yql/providers/yt/provider/yql_yt_ytflow_integration.cpp +++ b/yt/yql/providers/yt/provider/yql_yt_ytflow_integration.cpp @@ -289,6 +289,15 @@ public: sinkSettings.SetDoesExist(tableDesc.Meta->DoesExist); sinkSettings.SetTruncate(tableDesc.Intents & TYtTableIntent::Override); + + const auto& rowSpec = tableDesc.RowSpec; + if (rowSpec) { + for (const auto& [column, _] : rowSpec->GetForeignSort()) { + if (!rowSpec->ExpressionColumns.contains(column)) { + sinkSettings.AddKeyColumns(column); + } + } + } } settings.PackFrom(sinkSettings); diff --git a/yt/yql/providers/ytflow/integration/proto/yt.proto b/yt/yql/providers/ytflow/integration/proto/yt.proto index 2959ca4b6c8..4ebfad1059b 100644 --- a/yt/yql/providers/ytflow/integration/proto/yt.proto +++ b/yt/yql/providers/ytflow/integration/proto/yt.proto @@ -14,4 +14,5 @@ message TQYTSinkMessage { optional bytes RowType = 3; optional bool DoesExist = 5; optional bool Truncate = 6; + repeated string KeyColumns = 7; } diff --git a/yt/yql/tests/sql/suites/ytflow/multiple_sorted_outputs_with_different_keys.cfg b/yt/yql/tests/sql/suites/ytflow/multiple_sorted_outputs_with_different_keys.cfg new file mode 100644 index 00000000000..0e3721e4b77 --- /dev/null +++ b/yt/yql/tests/sql/suites/ytflow/multiple_sorted_outputs_with_different_keys.cfg @@ -0,0 +1,4 @@ +in Input input.txt +out OutputKey output_sorted_by_key.txt +out OutputValue output_sorted_by_value.txt +providers ytflow diff --git a/yt/yql/tests/sql/suites/ytflow/multiple_sorted_outputs_with_different_keys.yql b/yt/yql/tests/sql/suites/ytflow/multiple_sorted_outputs_with_different_keys.yql new file mode 100644 index 00000000000..b938bbf058a --- /dev/null +++ b/yt/yql/tests/sql/suites/ytflow/multiple_sorted_outputs_with_different_keys.yql @@ -0,0 +1,27 @@ +/* sort outputs */ + +use plato; + +pragma Engine = "ytflow"; + +pragma Ytflow.Cluster = "plato"; +pragma Ytflow.PipelinePath = "pipelines/test"; + +$lambda = ($row) -> { + $output_row_type = Struct<"key":Optional<String>, "value":Optional<Int64>>; + $variant_type = Variant<$output_row_type, $output_row_type>; + + return If( + $row.int64_field < 10, + Variant(<|"key":$row.string_field, "value":$row.int64_field|>, "0", $variant_type), + Variant(<|"key":$row.string_field, "value":$row.int64_field|>, "1", $variant_type) + ); +}; + +$sorted_by_key, $sorted_by_value = process Input using $lambda(TableRow()); + +insert into OutputKey +select * from $sorted_by_key; + +insert into OutputValue +select * from $sorted_by_value; diff --git a/yt/yql/tests/sql/suites/ytflow/output_composite_key_value_table.txt b/yt/yql/tests/sql/suites/ytflow/output_composite_key_value_table.txt new file mode 100644 index 00000000000..e69de29bb2d --- /dev/null +++ b/yt/yql/tests/sql/suites/ytflow/output_composite_key_value_table.txt diff --git a/yt/yql/tests/sql/suites/ytflow/output_composite_key_value_table.txt.attr b/yt/yql/tests/sql/suites/ytflow/output_composite_key_value_table.txt.attr new file mode 100644 index 00000000000..8f69a0c4304 --- /dev/null +++ b/yt/yql/tests/sql/suites/ytflow/output_composite_key_value_table.txt.attr @@ -0,0 +1,23 @@ +{ + "schema" = < + unique_keys = %true; + strict = %true; + > [ + { + "name" = "key_b"; + "type" = "string"; + "sort_order" = "ascending"; + }; + { + "name" = "key_a"; + "type" = "int64"; + "sort_order" = "ascending"; + }; + { + "name" = "kv_value"; + "type" = "int64"; + }; + ]; + "_yql_dynamic" = %true; + "_yql_dynamic_native_read" = %true; +} diff --git a/yt/yql/tests/sql/suites/ytflow/output_sorted_by_key.txt b/yt/yql/tests/sql/suites/ytflow/output_sorted_by_key.txt new file mode 100644 index 00000000000..e69de29bb2d --- /dev/null +++ b/yt/yql/tests/sql/suites/ytflow/output_sorted_by_key.txt diff --git a/yt/yql/tests/sql/suites/ytflow/output_sorted_by_key.txt.attr b/yt/yql/tests/sql/suites/ytflow/output_sorted_by_key.txt.attr new file mode 100644 index 00000000000..6086a84dbc9 --- /dev/null +++ b/yt/yql/tests/sql/suites/ytflow/output_sorted_by_key.txt.attr @@ -0,0 +1,18 @@ +{ + "schema" = < + unique_keys = %true; + strict = %true; + > [ + { + "name" = "key"; + "type" = "string"; + "sort_order" = "ascending"; + }; + { + "name" = "value"; + "type" = "int64"; + }; + ]; + "_yql_dynamic" = %true; + "_yql_dynamic_native_read" = %true; +} diff --git a/yt/yql/tests/sql/suites/ytflow/output_sorted_by_value.txt b/yt/yql/tests/sql/suites/ytflow/output_sorted_by_value.txt new file mode 100644 index 00000000000..e69de29bb2d --- /dev/null +++ b/yt/yql/tests/sql/suites/ytflow/output_sorted_by_value.txt diff --git a/yt/yql/tests/sql/suites/ytflow/output_sorted_by_value.txt.attr b/yt/yql/tests/sql/suites/ytflow/output_sorted_by_value.txt.attr new file mode 100644 index 00000000000..3acb927d6d8 --- /dev/null +++ b/yt/yql/tests/sql/suites/ytflow/output_sorted_by_value.txt.attr @@ -0,0 +1,18 @@ +{ + "schema" = < + unique_keys = %true; + strict = %true; + > [ + { + "name" = "value"; + "type" = "int64"; + "sort_order" = "ascending"; + }; + { + "name" = "key"; + "type" = "string"; + }; + ]; + "_yql_dynamic" = %true; + "_yql_dynamic_native_read" = %true; +} diff --git a/yt/yql/tests/sql/suites/ytflow/reuse_output_to_sorted_tables_with_different_keys.cfg b/yt/yql/tests/sql/suites/ytflow/reuse_output_to_sorted_tables_with_different_keys.cfg new file mode 100644 index 00000000000..0e3721e4b77 --- /dev/null +++ b/yt/yql/tests/sql/suites/ytflow/reuse_output_to_sorted_tables_with_different_keys.cfg @@ -0,0 +1,4 @@ +in Input input.txt +out OutputKey output_sorted_by_key.txt +out OutputValue output_sorted_by_value.txt +providers ytflow diff --git a/yt/yql/tests/sql/suites/ytflow/reuse_output_to_sorted_tables_with_different_keys.yql b/yt/yql/tests/sql/suites/ytflow/reuse_output_to_sorted_tables_with_different_keys.yql new file mode 100644 index 00000000000..7131563360a --- /dev/null +++ b/yt/yql/tests/sql/suites/ytflow/reuse_output_to_sorted_tables_with_different_keys.yql @@ -0,0 +1,16 @@ +/* sort outputs */ + +use plato; + +pragma Engine = "ytflow"; + +pragma Ytflow.Cluster = "plato"; +pragma Ytflow.PipelinePath = "pipelines/test"; + +$result_stream = select string_field as key, int64_field as value from Input; + +insert into OutputKey +select * from $result_stream; + +insert into OutputValue +select * from $result_stream; diff --git a/yt/yql/tests/sql/suites/ytflow/reuse_output_to_sorted_tables_with_same_key.cfg b/yt/yql/tests/sql/suites/ytflow/reuse_output_to_sorted_tables_with_same_key.cfg new file mode 100644 index 00000000000..596696392fb --- /dev/null +++ b/yt/yql/tests/sql/suites/ytflow/reuse_output_to_sorted_tables_with_same_key.cfg @@ -0,0 +1,4 @@ +in Input input.txt +out OutputFirst output_sorted_by_key.txt +out OutputSecond output_sorted_by_key.txt +providers ytflow diff --git a/yt/yql/tests/sql/suites/ytflow/reuse_output_to_sorted_tables_with_same_key.yql b/yt/yql/tests/sql/suites/ytflow/reuse_output_to_sorted_tables_with_same_key.yql new file mode 100644 index 00000000000..784e88094d3 --- /dev/null +++ b/yt/yql/tests/sql/suites/ytflow/reuse_output_to_sorted_tables_with_same_key.yql @@ -0,0 +1,16 @@ +/* sort outputs */ + +use plato; + +pragma Engine = "ytflow"; + +pragma Ytflow.Cluster = "plato"; +pragma Ytflow.PipelinePath = "pipelines/test"; + +$result_stream = select string_field as key, int64_field as value from Input; + +insert into OutputFirst +select * from $result_stream; + +insert into OutputSecond +select * from $result_stream; diff --git a/yt/yql/tests/sql/suites/ytflow/sorted_and_unsorted_outputs.cfg b/yt/yql/tests/sql/suites/ytflow/sorted_and_unsorted_outputs.cfg new file mode 100644 index 00000000000..d28699e18b4 --- /dev/null +++ b/yt/yql/tests/sql/suites/ytflow/sorted_and_unsorted_outputs.cfg @@ -0,0 +1,4 @@ +in Input input.txt +out OutputSorted output_sorted_by_key.txt +out OutputUnsorted output_value_multiple_outputs.txt +providers ytflow diff --git a/yt/yql/tests/sql/suites/ytflow/sorted_and_unsorted_outputs.yql b/yt/yql/tests/sql/suites/ytflow/sorted_and_unsorted_outputs.yql new file mode 100644 index 00000000000..b0ab691b823 --- /dev/null +++ b/yt/yql/tests/sql/suites/ytflow/sorted_and_unsorted_outputs.yql @@ -0,0 +1,28 @@ +/* sort outputs */ + +use plato; + +pragma Engine = "ytflow"; + +pragma Ytflow.Cluster = "plato"; +pragma Ytflow.PipelinePath = "pipelines/test"; + +$lambda = ($row) -> { + $sorted_row_type = Struct<"key":Optional<String>, "value":Optional<Int64>>; + $unsorted_row_type = Struct<"value":Optional<Int64>>; + $variant_type = Variant<$sorted_row_type, $unsorted_row_type>; + + return If( + $row.int64_field < 10, + Variant(<|"key":$row.string_field, "value":$row.int64_field|>, "0", $variant_type), + Variant(<|"value":$row.int64_field|>, "1", $variant_type) + ); +}; + +$sorted, $unsorted = process Input using $lambda(TableRow()); + +insert into OutputSorted +select * from $sorted; + +insert into OutputUnsorted +select * from $unsorted; diff --git a/yt/yql/tests/sql/suites/ytflow/write_to_sorted_table_with_composite_key.cfg b/yt/yql/tests/sql/suites/ytflow/write_to_sorted_table_with_composite_key.cfg new file mode 100644 index 00000000000..ae0ae973f2e --- /dev/null +++ b/yt/yql/tests/sql/suites/ytflow/write_to_sorted_table_with_composite_key.cfg @@ -0,0 +1,3 @@ +in Input input_stream_composite_key.txt +out Output output_composite_key_value_table.txt +providers ytflow diff --git a/yt/yql/tests/sql/suites/ytflow/write_to_sorted_table_with_composite_key.yql b/yt/yql/tests/sql/suites/ytflow/write_to_sorted_table_with_composite_key.yql new file mode 100644 index 00000000000..37f5aa55876 --- /dev/null +++ b/yt/yql/tests/sql/suites/ytflow/write_to_sorted_table_with_composite_key.yql @@ -0,0 +1,15 @@ +/* sort outputs */ + +use plato; + +pragma Engine = "ytflow"; + +pragma Ytflow.Cluster = "plato"; +pragma Ytflow.PipelinePath = "pipelines/test"; + +insert into Output +select + key_a, + key_b || "_b" as key_b, + value as kv_value +from Input; diff --git a/yt/yql/tools/mrjob/mrjob.cpp b/yt/yql/tools/mrjob/mrjob.cpp index 52611c7abf9..d8a8c058740 100644 --- a/yt/yql/tools/mrjob/mrjob.cpp +++ b/yt/yql/tools/mrjob/mrjob.cpp @@ -5,17 +5,11 @@ #include <yt/cpp/mapreduce/client/init.h> -#include <tcmalloc/malloc_extension.h> - #include <util/system/env.h> #include <util/system/mlock.h> #include <util/system/yassert.h> int main(int argc, const char *argv[]) { - // disable tcmalloc sampling - tcmalloc::MallocExtension::SetGuardedSamplingInterval(-1); - tcmalloc::MallocExtension::SetProfileSamplingInterval(0); - Y_UNUSED(NKikimr::NUdf::GetStaticSymbols()); try { LockAllMemory(LockCurrentMemory | LockFutureMemory); diff --git a/yt/yql/tools/mrjob/ya.make b/yt/yql/tools/mrjob/ya.make index 84ba0a368b0..4d88e6b5b59 100644 --- a/yt/yql/tools/mrjob/ya.make +++ b/yt/yql/tools/mrjob/ya.make @@ -1,10 +1,6 @@ PROGRAM(mrjob) -IF (OS_LINUX) - ALLOCATOR(TCMALLOC) -ELSE() - ALLOCATOR(J) -ENDIF() +ALLOCATOR(J) SRCS( mrjob.cpp @@ -17,7 +13,6 @@ IF (OS_LINUX) ENDIF() PEERDIR( - contrib/libs/tcmalloc/malloc_extension yt/cpp/mapreduce/client yql/essentials/public/udf/service/terminate_policy yql/essentials/providers/common/gateway diff --git a/yt/yt/RpcProxyProtocolVersion.txt b/yt/yt/RpcProxyProtocolVersion.txt index 9eea6d6e60a..c1cd2292c32 100644 --- a/yt/yt/RpcProxyProtocolVersion.txt +++ b/yt/yt/RpcProxyProtocolVersion.txt @@ -16,7 +16,7 @@ SET(YT_RPC_PROXY_SERVER_PROTOCOL_VERSION_MINOR 2) SET(YT_RPC_MODIFY_ROWS_STRONG_LOCKS_VERSION 2) # RPC python client has protocol version plus some patch version -SET(YT_RPC_PYTHON_BINDINGS_PATCH_VERSION 22) +SET(YT_RPC_PYTHON_BINDINGS_PATCH_VERSION 23) # YT proto package has protocol version plus some patch version SET(YT_PROTO_PACKAGE_PATCH_VERSION 15) diff --git a/yt/yt/client/api/rpc_proxy/client_impl.h b/yt/yt/client/api/rpc_proxy/client_impl.h index fa7dfdb1a26..51905bdf769 100644 --- a/yt/yt/client/api/rpc_proxy/client_impl.h +++ b/yt/yt/client/api/rpc_proxy/client_impl.h @@ -705,7 +705,6 @@ public: private: const TConnectionPtr Connection_; const TClientOptions ClientOptions_; - const NRpc::TDynamicChannelPoolPtr ChannelPool_; const NRpc::IChannelPtr RetryingChannel_; TLazyIntrusivePtr<NTabletClient::ITableMountCache> TableMountCache_; diff --git a/yt/yt/client/object_client/helpers.cpp b/yt/yt/client/object_client/helpers.cpp index facf1ae54d8..45746d9150e 100644 --- a/yt/yt/client/object_client/helpers.cpp +++ b/yt/yt/client/object_client/helpers.cpp @@ -140,6 +140,7 @@ bool IsVersionedType(EObjectType type) case EObjectType::Pipeline: case EObjectType::QueueConsumer: case EObjectType::QueueProducer: + case EObjectType::QueueMultiConsumer: case EObjectType::CypressProxyMap: return true; diff --git a/yt/yt/client/object_client/public.h b/yt/yt/client/object_client/public.h index ae8d918bab1..6c6e9a31fba 100644 --- a/yt/yt/client/object_client/public.h +++ b/yt/yt/client/object_client/public.h @@ -4,14 +4,14 @@ #include <yt/yt/client/job_tracker_client/public.h> +#include <library/cpp/yt/compact_containers/compact_flat_set.h> +#include <library/cpp/yt/compact_containers/compact_vector.h> + #include <library/cpp/yt/misc/enum.h> #include <library/cpp/yt/misc/guid.h> #include <library/cpp/yt/misc/hash.h> #include <library/cpp/yt/misc/strong_typedef.h> -#include <library/cpp/yt/compact_containers/compact_vector.h> -#include <library/cpp/yt/compact_containers/compact_flat_set.h> - #include <library/cpp/yt/string/string_builder.h> namespace NYT::NObjectClient { @@ -368,6 +368,7 @@ DEFINE_ENUM(EObjectType, // Queue stuff ((QueueConsumer) (1700)) ((QueueProducer) (1701)) + ((QueueMultiConsumer) (1702)) ); //! A bit mask marking schema types. diff --git a/yt/yt/client/sequoia_client/public.h b/yt/yt/client/sequoia_client/public.h index de32a1a1b9c..b9c79e97750 100644 --- a/yt/yt/client/sequoia_client/public.h +++ b/yt/yt/client/sequoia_client/public.h @@ -12,6 +12,7 @@ YT_DEFINE_ERROR_ENUM( ((SequoiaRetriableError) (6002)) ((InvalidSequoiaReign) (6003)) ((TransactionActionFailedOnMasterCell) (6004)) + ((SequoiaSubtreeTooLarge) (6005)) ); //////////////////////////////////////////////////////////////////////////////// diff --git a/yt/yt/core/logging/config.cpp b/yt/yt/core/logging/config.cpp index e5ea991aa72..f21f8aa336f 100644 --- a/yt/yt/core/logging/config.cpp +++ b/yt/yt/core/logging/config.cpp @@ -57,6 +57,8 @@ void TLogWriterConfig::Register(TRegistrar registrar) .Default(true); registrar.Parameter("enable_host_field", &TThis::EnableHostField) .Default(false); + registrar.Parameter("enable_native_tags", &TThis::EnableNativeTags) + .Default(false); registrar.Parameter("json_format", &TThis::JsonFormat) .Default(); registrar.Parameter("yson_format", &TThis::YsonFormat) diff --git a/yt/yt/core/logging/config.h b/yt/yt/core/logging/config.h index 452f41b4233..d0b1cbf366b 100644 --- a/yt/yt/core/logging/config.h +++ b/yt/yt/core/logging/config.h @@ -10,6 +10,8 @@ #include <yt/yt/core/ytree/public.h> #include <yt/yt/core/ytree/yson_struct.h> +#include <library/cpp/yt/yson_string/string.h> + namespace NYT::NLogging { //////////////////////////////////////////////////////////////////////////////// @@ -35,9 +37,15 @@ struct TLogWriterConfig bool EnableSourceLocation; //! Structured formatter options. + //! If set, every event carries the system fields "instant", "level" and "category". bool EnableSystemFields; + //! If set, every event carries the "host" field with the local hostname. bool EnableHostField; - THashMap<std::string, NYTree::INodePtr> CommonFields; + //! If set, a tagged message's tags go to a nested "tags" map; otherwise they + //! are folded into the "message" field as |Message (Key: Value, ...)|. + bool EnableNativeTags; + //! Fields injected into every event as |key -> raw YSON value|. + THashMap<std::string, NYson::TYsonString> CommonFields; NJson::TJsonFormatConfigPtr JsonFormat; NYson::EYsonFormat YsonFormat; diff --git a/yt/yt/core/logging/formatter.cpp b/yt/yt/core/logging/formatter.cpp index 318b4f72653..1dd37fba492 100644 --- a/yt/yt/core/logging/formatter.cpp +++ b/yt/yt/core/logging/formatter.cpp @@ -2,6 +2,9 @@ #include "private.h" +#include <library/cpp/yt/logging/tagged_payload.h> +#include <library/cpp/yt/logging/structured_payload.h> + #include <yt/yt/core/json/json_writer.h> #include <yt/yt/core/ytree/fluent.h> @@ -18,9 +21,9 @@ using namespace NYson; //////////////////////////////////////////////////////////////////////////////// -TPlainTextLogFormatter::TPlainTextLogFormatter(bool enableSourceLocation) - : TLogFormatterBase(enableSourceLocation) - , EventFormatter_(enableSourceLocation) +TPlainTextLogFormatter::TPlainTextLogFormatter(TPlainTextLogFormatterOptions options) + : TLogFormatterBase(TLogFormatterBaseOptions{.EnableSourceLocation = options.EnableSourceLocation}) + , EventFormatter_(options.EnableSourceLocation) { } i64 TPlainTextLogFormatter::WriteFormatted(IOutputStream* outputStream, const TLogEvent& event) @@ -45,34 +48,32 @@ void TPlainTextLogFormatter::WriteLogReopenSeparator(IOutputStream* outputStream //////////////////////////////////////////////////////////////////////////////// -TLogFormatterBase::TLogFormatterBase(bool enableSourceLocation) - : EnableSourceLocation_(enableSourceLocation) +TLogFormatterBase::TLogFormatterBase(TLogFormatterBaseOptions options) + : Options_(std::move(options)) { } bool TLogFormatterBase::IsSourceLocationEnabled() const { - return EnableSourceLocation_; + return Options_.EnableSourceLocation; } //////////////////////////////////////////////////////////////////////////////// -TStructuredLogFormatter::TStructuredLogFormatter( - ELogFormat format, - THashMap<std::string, INodePtr> commonFields, - bool enableSourceLocation, - bool enableSystemFields, - bool enableHostField, - NJson::TJsonFormatConfigPtr jsonFormat, - EYsonFormat ysonFormat) - : TLogFormatterBase(enableSourceLocation) - , Format_(format) - , CommonFields_(std::move(commonFields)) - , EnableSystemFields_(enableSystemFields) - , EnableHostField_(enableHostField) - , JsonFormat_(!jsonFormat && (Format_ == ELogFormat::Json) - ? New<NJson::TJsonFormatConfig>() - : std::move(jsonFormat)) - , YsonFormat_(ysonFormat) +namespace { + +TStructuredLogFormatterOptions NormalizeStructuredOptions(TStructuredLogFormatterOptions options) +{ + if (!options.JsonFormat && options.Format == ELogFormat::Json) { + options.JsonFormat = New<NJson::TJsonFormatConfig>(); + } + return options; +} + +} // namespace + +TStructuredLogFormatter::TStructuredLogFormatter(TStructuredLogFormatterOptions options) + : TLogFormatterBase(TLogFormatterBaseOptions{.EnableSourceLocation = options.EnableSourceLocation}) + , Options_(NormalizeStructuredOptions(std::move(options))) { } i64 TStructuredLogFormatter::WriteFormatted(IOutputStream* stream, const TLogEvent& event) @@ -84,17 +85,17 @@ i64 TStructuredLogFormatter::WriteFormatted(IOutputStream* stream, const TLogEve auto countingStream = TCountingOutput(stream); std::unique_ptr<IFlushableYsonConsumer> consumer; - switch (Format_) { + switch (Options_.Format) { case ELogFormat::Json: - YT_VERIFY(JsonFormat_); - consumer = NJson::CreateJsonConsumer(&countingStream, EYsonType::Node, JsonFormat_); + YT_VERIFY(Options_.JsonFormat); + consumer = NJson::CreateJsonConsumer(&countingStream, EYsonType::Node, Options_.JsonFormat); break; case ELogFormat::Yson: consumer = std::make_unique<TYsonWriter>( &countingStream, - YsonFormat_, + Options_.YsonFormat, EYsonType::Node, - /*enableRaw*/ YsonFormat_ == EYsonFormat::Binary); + /*enableRaw*/ Options_.YsonFormat == EYsonFormat::Binary); break; default: YT_ABORT(); @@ -105,23 +106,40 @@ i64 TStructuredLogFormatter::WriteFormatted(IOutputStream* stream, const TLogEve BuildYsonFluently(consumer.get()) .BeginMap() - .DoFor(CommonFields_, [] (auto fluent, auto item) { + .DoFor(Options_.CommonFields, [] (auto fluent, auto item) { fluent.Item(item.first).Value(item.second); }) - .DoIf(event.MessageKind == ELogMessageKind::Structured, [&] (auto fluent) { - fluent.Items(TYsonString(event.MessageRef, EYsonType::MapFragment)); + .DoIf(std::holds_alternative<TStructuredLogEventPayload>(event.Payload), [&] (auto fluent) { + fluent.Items(TYsonString(GetYsonFromStructuredPayload(std::get<TStructuredLogEventPayload>(event.Payload)))); }) - .DoIf(event.MessageKind == ELogMessageKind::Unstructured, [&] (auto fluent) { - fluent.Item("message").Value(event.MessageRef.ToStringBuf()); + .DoIf(std::holds_alternative<TTaggedLogEventPayload>(event.Payload), [&] (auto fluent) { + const auto& taggedPayload = std::get<TTaggedLogEventPayload>(event.Payload); + if (Options_.EnableNativeTags) { + TTaggedPayloadReader reader(taggedPayload); + auto message = reader.ReadMessage(); + auto firstTag = reader.TryReadTag(); + fluent + .Item("message").Value(message) + .DoIf(firstTag.has_value(), [&] (auto fluent) { + fluent.Item("tags").DoMap([&] (auto fluent) { + for (auto tag = firstTag; tag; tag = reader.TryReadTag()) { + fluent.Item(tag->Key).Value(tag->Value); + } + }); + }); + } else { + // Fold the tags into the message: |Message (Key: Value, ...)|. + fluent.Item("message").Value(FormatTaggedPayload(taggedPayload)); + } }) - .DoIf(EnableSystemFields_, [&] (auto fluent) { + .DoIf(Options_.EnableSystemFields, [&] (auto fluent) { fluent .Item("instant").Value(dateTimeBuffer.GetBuffer()) .Item("level").Value(FormatEnum(event.Level)) .Item("category").Value(event.Category->Name); }) // TODO(achulkov2): The presence of different system fields should be controller by a flag enum instead of multiple boolean options. - .DoIf(EnableHostField_, [&] (auto fluent) { + .DoIf(Options_.EnableHostField, [&] (auto fluent) { fluent.Item("host").Value(NNet::GetLocalHostName()); }) .DoIf(event.Family == ELogFamily::PlainText, [&] (auto fluent) { @@ -139,7 +157,7 @@ i64 TStructuredLogFormatter::WriteFormatted(IOutputStream* stream, const TLogEve .EndMap(); consumer->Flush(); - if (Format_ == ELogFormat::Yson) { + if (Options_.Format == ELogFormat::Yson) { // In order to obtain proper list fragment, we must manually insert trailing semicolon in each line. countingStream.Write(';'); } diff --git a/yt/yt/core/logging/formatter.h b/yt/yt/core/logging/formatter.h index df7c1a46acd..91d0569484b 100644 --- a/yt/yt/core/logging/formatter.h +++ b/yt/yt/core/logging/formatter.h @@ -6,6 +6,8 @@ #include <library/cpp/yt/logging/plain_text_formatter/formatter.h> +#include <library/cpp/yt/yson_string/string.h> + namespace NYT::NLogging { //////////////////////////////////////////////////////////////////////////////// @@ -20,25 +22,35 @@ struct ILogFormatter //////////////////////////////////////////////////////////////////////////////// +struct TLogFormatterBaseOptions +{ + bool EnableSourceLocation = false; +}; + class TLogFormatterBase : public ILogFormatter { protected: - TLogFormatterBase(bool enableSourceLocation); + explicit TLogFormatterBase(TLogFormatterBaseOptions options); bool IsSourceLocationEnabled() const; private: - const bool EnableSourceLocation_; + const TLogFormatterBaseOptions Options_; }; //////////////////////////////////////////////////////////////////////////////// +struct TPlainTextLogFormatterOptions +{ + bool EnableSourceLocation = false; +}; + class TPlainTextLogFormatter : public TLogFormatterBase { public: - explicit TPlainTextLogFormatter(bool enableSourceLocation = false); + explicit TPlainTextLogFormatter(TPlainTextLogFormatterOptions options = {}); i64 WriteFormatted(IOutputStream* outputStream, const TLogEvent& event) override; void WriteLogReopenSeparator(IOutputStream* outputStream) override; @@ -50,29 +62,30 @@ private: //////////////////////////////////////////////////////////////////////////////// +struct TStructuredLogFormatterOptions +{ + ELogFormat Format = ELogFormat::Json; + //! Fields injected into every event as |key -> raw YSON value|. + THashMap<std::string, NYson::TYsonString> CommonFields; + bool EnableSourceLocation = false; + bool EnableSystemFields = true; + bool EnableHostField = false; + bool EnableNativeTags = false; + NJson::TJsonFormatConfigPtr JsonFormat; + NYson::EYsonFormat YsonFormat = NYson::EYsonFormat::Text; +}; + class TStructuredLogFormatter : public TLogFormatterBase { public: - TStructuredLogFormatter( - ELogFormat format, - THashMap<std::string, NYTree::INodePtr> commonFields, - bool enableSourceLocation = false, - bool enableSystemFields = true, - bool enableHostField = false, - NJson::TJsonFormatConfigPtr jsonFormat = nullptr, - NYson::EYsonFormat ysonFormat = NYson::EYsonFormat::Text); + explicit TStructuredLogFormatter(TStructuredLogFormatterOptions options); i64 WriteFormatted(IOutputStream* outputStream, const TLogEvent& event) override; void WriteLogReopenSeparator(IOutputStream* outputStream) override; private: - const ELogFormat Format_; - const THashMap<std::string, NYTree::INodePtr> CommonFields_; - const bool EnableSystemFields_; - const bool EnableHostField_; - const NJson::TJsonFormatConfigPtr JsonFormat_; - const NYson::EYsonFormat YsonFormat_; + const TStructuredLogFormatterOptions Options_; TCachingDateFormatter CachingDateFormatter_; }; diff --git a/yt/yt/core/logging/log-inl.h b/yt/yt/core/logging/log-inl.h index 52178a5ab75..076d3c28d32 100644 --- a/yt/yt/core/logging/log-inl.h +++ b/yt/yt/core/logging/log-inl.h @@ -20,11 +20,13 @@ TLogMessage BuildLogMessage( TFormatString<TArgs...> format, TArgs&&... args) { - TMessageStringBuilder builder; - AppendLogMessageWithFormat(&builder, loggingContext, logger, format.Get(), std::forward<TArgs>(args)...); - builder.AppendChar('\n'); - FormatValue(&builder, error, TStringBuf("v")); - return {builder.Flush(), format.Get()}; + TTaggedPayloadWriter writer; + auto* builder = writer.BeginMessage(); + AppendLogMessageWithFormat(builder, loggingContext, logger, format.Get(), std::forward<TArgs>(args)...); + builder->AppendChar('\n'); + FormatValue(builder, error, TStringBuf("v")); + writer.EndMessage(); + return {writer.Finish(), format.Get()}; } } // namespace NDetail diff --git a/yt/yt/core/logging/log_manager.cpp b/yt/yt/core/logging/log_manager.cpp index bda678f2d63..aa41d1fb4e0 100644 --- a/yt/yt/core/logging/log_manager.cpp +++ b/yt/yt/core/logging/log_manager.cpp @@ -10,6 +10,9 @@ #include "stream_log_writer.h" #include "system_log_event_provider.h" +#include <library/cpp/yt/logging/private.h> +#include <library/cpp/yt/logging/tagged_payload.h> + #include <yt/yt/core/concurrency/helpers.h> #include <yt/yt/core/concurrency/periodic_executor.h> #include <yt/yt/core/concurrency/scheduler_thread.h> @@ -417,7 +420,7 @@ public: // Limit message length. TRawFormatter<1024> truncatedMessage; - truncatedMessage.AppendString(event.MessageRef.ToStringBuf()); + truncatedMessage.AppendString(GetMessageFromTaggedPayload(std::get<TTaggedLogEventPayload>(event.Payload))); // NB(coteeq): It's safe to save event.SourceFile since it was // generated by __LOCATION__ macro. @@ -701,19 +704,22 @@ private: switch (writerConfig->Format) { case ELogFormat::PlainText: - return std::make_unique<TPlainTextLogFormatter>( - writerConfig->EnableSourceLocation); + return std::make_unique<TPlainTextLogFormatter>(TPlainTextLogFormatterOptions{ + .EnableSourceLocation = writerConfig->EnableSourceLocation, + }); case ELogFormat::Json: [[fallthrough]]; case ELogFormat::Yson: - return std::make_unique<TStructuredLogFormatter>( - writerConfig->Format, - writerConfig->CommonFields, - writerConfig->EnableSourceLocation, - writerConfig->EnableSystemFields, - writerConfig->EnableHostField, - writerConfig->JsonFormat, - writerConfig->YsonFormat); + return std::make_unique<TStructuredLogFormatter>(TStructuredLogFormatterOptions{ + .Format = writerConfig->Format, + .CommonFields = writerConfig->CommonFields, + .EnableSourceLocation = writerConfig->EnableSourceLocation, + .EnableSystemFields = writerConfig->EnableSystemFields, + .EnableHostField = writerConfig->EnableHostField, + .EnableNativeTags = writerConfig->EnableNativeTags, + .JsonFormat = writerConfig->JsonFormat, + .YsonFormat = writerConfig->YsonFormat, + }); default: YT_ABORT(); @@ -815,7 +821,9 @@ private: if (event.Anchor) { event.Anchor->MessageCounter.Current += 1; - event.Anchor->ByteCounter.Current += std::ssize(event.MessageRef); + event.Anchor->ByteCounter.Current += std::visit( + [] (const auto& payload) { return std::ssize(payload.Underlying()); }, + event.Payload); } for (const auto& writer : GetWriters(config, event)) { diff --git a/yt/yt/core/logging/system_log_event_provider.cpp b/yt/yt/core/logging/system_log_event_provider.cpp index 9bec1d61515..bef41ce1c76 100644 --- a/yt/yt/core/logging/system_log_event_provider.cpp +++ b/yt/yt/core/logging/system_log_event_provider.cpp @@ -6,6 +6,8 @@ #include <yt/yt/core/ytree/fluent.h> #include <library/cpp/yt/logging/logger.h> +#include <library/cpp/yt/logging/tagged_payload.h> +#include <library/cpp/yt/logging/structured_payload.h> namespace NYT::NLogging { @@ -47,8 +49,7 @@ public: return TLogEvent{ .Category = Logger().GetCategory(), .Level = ELogLevel::Info, - .MessageKind = ELogMessageKind::Unstructured, - .MessageRef = TSharedRef::FromString(Format("Logging started (Version: %v, BuildHost: %v, BuildTime: %v)", + .Payload = MakeTaggedPayloadFromMessage(Format("Logging started (Version: %v, BuildHost: %v, BuildTime: %v)", GetVersion(), GetBuildHost(), GetBuildTime())), @@ -65,8 +66,7 @@ public: return TLogEvent{ .Category = Logger().GetCategory(), .Level = ELogLevel::Info, - .MessageKind = ELogMessageKind::Unstructured, - .MessageRef = TSharedRef::FromString(Format("Skipped log records in last second (Count: %v, SkippedBy: %v)", + .Payload = MakeTaggedPayloadFromMessage(Format("Skipped log records in last second (Count: %v, SkippedBy: %v)", count, skippedBy)), .Instant = GetCpuInstant(), @@ -89,14 +89,12 @@ public: return TLogEvent{ .Category = Logger().GetCategory(), .Level = ELogLevel::Info, - .MessageKind = ELogMessageKind::Structured, - .MessageRef = BuildYsonStringFluently<NYson::EYsonType::MapFragment>() + .Payload = MakeStructuredPayloadFromYson(BuildYsonStringFluently<NYson::EYsonType::MapFragment>() .Item("message").Value("Logging started") .Item("version").Value(GetVersion()) .Item("build_host").Value(GetBuildHost()) .Item("build_time").Value(GetBuildTime()) - .Finish() - .ToSharedRef(), + .Finish()), .Instant = GetCpuInstant(), }; } @@ -110,13 +108,11 @@ public: return TLogEvent{ .Category = Logger().GetCategory(), .Level = ELogLevel::Info, - .MessageKind = ELogMessageKind::Structured, - .MessageRef = BuildYsonStringFluently<NYson::EYsonType::MapFragment>() + .Payload = MakeStructuredPayloadFromYson(BuildYsonStringFluently<NYson::EYsonType::MapFragment>() .Item("message").Value("Events skipped") .Item("skipped_by").Value(skippedBy) .Item("events_skipped").Value(count) - .Finish() - .ToSharedRef(), + .Finish()), .Instant = GetCpuInstant(), }; } diff --git a/yt/yt/core/logging/unittests/logging_ut.cpp b/yt/yt/core/logging/unittests/logging_ut.cpp index ce9c6c006aa..e4c03879b84 100644 --- a/yt/yt/core/logging/unittests/logging_ut.cpp +++ b/yt/yt/core/logging/unittests/logging_ut.cpp @@ -20,6 +20,9 @@ #include <yt/yt/core/logging/system_log_event_provider.h> #include <yt/yt/core/logging/zstd_log_codec.h> +#include <library/cpp/yt/logging/tagged_payload.h> +#include <library/cpp/yt/logging/structured_payload.h> + #include <yt/yt/core/misc/fs.h> #include <yt/yt/core/test_framework/framework.h> @@ -55,6 +58,8 @@ namespace NYT::NLogging { namespace { +using namespace NDetail; + using namespace NYTree; using namespace NConcurrency; using namespace NYson; @@ -63,6 +68,26 @@ using namespace NCoreDump; //////////////////////////////////////////////////////////////////////////////// +void WriteMessage(TTaggedPayloadWriter* writer, TStringBuf message) +{ + writer->BeginMessage()->AppendString(message); + writer->EndMessage(); +} + +void WriteTag(TTaggedPayloadWriter* writer, TStringBuf key, TStringBuf value) +{ + writer->BeginTag(key)->AppendString(value); + writer->EndTag(); +} + +void WriteWellKnownTag(TTaggedPayloadWriter* writer, TStringBuf key, TStringBuf value) +{ + writer->BeginWellKnownTag(key)->AppendString(value); + writer->EndTag(); +} + +//////////////////////////////////////////////////////////////////////////////// + YT_DEFINE_GLOBAL(NLogging::TLogger, Logger, "Test"); std::string GenerateLogFileName() @@ -98,13 +123,13 @@ protected: void WritePlainTextEvent(const ILogWriterPtr& writer) { - TLogEvent event; - event.Family = ELogFamily::PlainText; - event.Category = Logger().GetCategory(); - event.Level = ELogLevel::Debug; - event.MessageRef = TSharedRef::FromString(std::string("message")); - event.MessageKind = ELogMessageKind::Unstructured; - event.ThreadId = 0xba; + TLogEvent event{ + .Category = Logger().GetCategory(), + .Level = ELogLevel::Debug, + .Family = ELogFamily::PlainText, + .Payload = MakeTaggedPayloadFromMessage("message"), + .ThreadId = 0xba, + }; WriteEvent(writer, event); } @@ -530,16 +555,16 @@ TEST_F(TLoggingTest, ThreadMinLogLevel) TEST_F(TLoggingTest, PlainTextLoggingStructuredFormatter) { - TLogEvent event; - event.Family = ELogFamily::PlainText; - event.Category = Logger().GetCategory(); - event.Level = ELogLevel::Debug; - event.MessageRef = TSharedRef::FromString(std::string("test_message")); - event.MessageKind = ELogMessageKind::Unstructured; - event.FiberId = 31; - event.TraceId = TGuid(1, 2, 3, 4); - event.SourceFile = "a/b.cpp"; - event.SourceLine = 123; + TLogEvent event{ + .Category = Logger().GetCategory(), + .Level = ELogLevel::Debug, + .Family = ELogFamily::PlainText, + .Payload = MakeTaggedPayloadFromMessage("test_message"), + .FiberId = 31, + .TraceId = TTraceId(1, 2, 3, 4), + .SourceFile = "a/b.cpp", + .SourceLine = 123, + }; for (auto enableSourceLocation : {false, true}) { for (auto format : {ELogFormat::Yson, ELogFormat::Json}) { @@ -549,7 +574,7 @@ TEST_F(TLoggingTest, PlainTextLoggingStructuredFormatter) writerConfig->FileName = logFile.Name(); auto writer = CreateFileLogWriter( - std::make_unique<TStructuredLogFormatter>(format, THashMap<std::string, INodePtr>{}, enableSourceLocation), + std::make_unique<TStructuredLogFormatter>(TStructuredLogFormatterOptions{.Format = format, .EnableSourceLocation = enableSourceLocation}), CreateDefaultSystemLogEventProvider(writerConfig), "test_writer", writerConfig, @@ -579,18 +604,16 @@ TEST_F(TLoggingTest, PlainTextLoggingStructuredFormatter) TEST_F(TLoggingTest, StructuredLogging) { - TLogEvent event; - event.Family = ELogFamily::Structured; - event.Category = Logger().GetCategory(); - event.Level = ELogLevel::Debug; - event.MessageRef = BuildYsonStringFluently<EYsonType::MapFragment>() - .Item("message").Value("test_message") - .Finish() - .ToSharedRef(); - event.MessageKind = ELogMessageKind::Structured; - - event.FiberId = 31; - event.TraceId = TGuid(1, 2, 3, 4); + TLogEvent event{ + .Category = Logger().GetCategory(), + .Level = ELogLevel::Debug, + .Family = ELogFamily::Structured, + .Payload = MakeStructuredPayloadFromYson(BuildYsonStringFluently<EYsonType::MapFragment>() + .Item("message").Value("test_message") + .Finish()), + .FiberId = 31, + .TraceId = TTraceId(1, 2, 3, 4), + }; for (auto format : {ELogFormat::Yson, ELogFormat::Json}) { TTempFile logFile(GenerateLogFileName()); @@ -599,7 +622,7 @@ TEST_F(TLoggingTest, StructuredLogging) writerConfig->FileName = logFile.Name(); auto writer = CreateFileLogWriter( - std::make_unique<TStructuredLogFormatter>(format, THashMap<std::string, INodePtr>{}), + std::make_unique<TStructuredLogFormatter>(TStructuredLogFormatterOptions{.Format = format}), CreateDefaultSystemLogEventProvider(writerConfig), "test_writer", writerConfig, @@ -728,14 +751,14 @@ TEST_F(TLoggingTest, SlogMacroConstexprLogging) } } -TEST_F(TLoggingTest, UnstructuredLogging) +TEST_F(TLoggingTest, TaggedLogging) { - TLogEvent event; - event.Family = ELogFamily::Structured; - event.Category = Logger().GetCategory(); - event.Level = ELogLevel::Debug; - event.MessageRef = TSharedRef::FromString(std::string("test_message")); - event.MessageKind = ELogMessageKind::Unstructured; + TLogEvent event{ + .Category = Logger().GetCategory(), + .Level = ELogLevel::Debug, + .Family = ELogFamily::Structured, + .Payload = MakeTaggedPayloadFromMessage("test_message"), + }; for (auto format : {ELogFormat::Yson, ELogFormat::Json}) { TTempFile logFile(GenerateLogFileName()); @@ -744,7 +767,7 @@ TEST_F(TLoggingTest, UnstructuredLogging) writerConfig->FileName = logFile.Name(); auto writer = CreateFileLogWriter( - std::make_unique<TStructuredLogFormatter>(format, THashMap<std::string, INodePtr>{}), + std::make_unique<TStructuredLogFormatter>(TStructuredLogFormatterOptions{.Format = format}), CreateDefaultSystemLogEventProvider(writerConfig), "test_writer", writerConfig, @@ -768,17 +791,16 @@ TEST_F(TLoggingTest, StructuredLoggingJsonFormat) std::string longString(1000, 'a'); std::string longStringPrefix(100, 'a'); - TLogEvent event; - event.Family = ELogFamily::Structured; - event.Category = Logger().GetCategory(); - event.Level = ELogLevel::Debug; - event.MessageRef = BuildYsonStringFluently<EYsonType::MapFragment>() - .Item("message").Value("test_message") - .Item("nan_value").Value(std::nan("1")) - .Item("long_string_value").Value(longString) - .Finish() - .ToSharedRef(); - event.MessageKind = ELogMessageKind::Structured; + TLogEvent event{ + .Category = Logger().GetCategory(), + .Level = ELogLevel::Debug, + .Family = ELogFamily::Structured, + .Payload = MakeStructuredPayloadFromYson(BuildYsonStringFluently<EYsonType::MapFragment>() + .Item("message").Value("test_message") + .Item("nan_value").Value(std::nan("1")) + .Item("long_string_value").Value(longString) + .Finish()), + }; auto jsonFormat = New<TJsonFormatConfig>(); jsonFormat->StringifyNanAndInfinity = true; @@ -789,13 +811,10 @@ TEST_F(TLoggingTest, StructuredLoggingJsonFormat) auto writerConfig = New<TFileLogWriterConfig>(); writerConfig->FileName = logFile.Name(); - auto formatter = std::make_unique<TStructuredLogFormatter>( - ELogFormat::Json, - /*commonFields*/ THashMap<std::string, INodePtr>{}, - /*enableSourceLocation*/ false, - /*enableSystemFields*/ true, - /*enableHostField*/ false, - jsonFormat); + auto formatter = std::make_unique<TStructuredLogFormatter>(TStructuredLogFormatterOptions{ + .Format = ELogFormat::Json, + .JsonFormat = jsonFormat, + }); auto writer = CreateFileLogWriter( std::move(formatter), @@ -943,22 +962,20 @@ TEST_F(TLoggingTest, StructuredValidationWithSamplingRate) TEST_F(TLoggingTest, StructuredLoggingDisableSystemFields) { - TLogEvent event; - event.Family = ELogFamily::Structured; - event.Category = Logger().GetCategory(); - event.Level = ELogLevel::Debug; - event.MessageRef = BuildYsonStringFluently<EYsonType::MapFragment>() - .Item("message").Value("test_message") - .Finish() - .ToSharedRef(); - event.MessageKind = ELogMessageKind::Structured; + TLogEvent event{ + .Category = Logger().GetCategory(), + .Level = ELogLevel::Debug, + .Family = ELogFamily::Structured, + .Payload = MakeStructuredPayloadFromYson(BuildYsonStringFluently<EYsonType::MapFragment>() + .Item("message").Value("test_message") + .Finish()), + }; - auto formatter = std::make_unique<TStructuredLogFormatter>( - ELogFormat::Yson, - /*commonFields*/ THashMap<std::string, INodePtr>{}, - /*enableControlMessages*/ true, - /*enableSourceLocation*/ false, - /*enableSystemFields*/ false); + auto formatter = std::make_unique<TStructuredLogFormatter>(TStructuredLogFormatterOptions{ + .Format = ELogFormat::Yson, + .EnableSourceLocation = false, + .EnableSystemFields = false, + }); TStringStream stringStream; formatter->WriteFormatted(&stringStream, event); @@ -971,6 +988,63 @@ TEST_F(TLoggingTest, StructuredLoggingDisableSystemFields) EXPECT_EQ(message->FindChild("category"), nullptr); } +TEST_F(TLoggingTest, PlainTextFormatterTags) +{ + TTaggedPayloadWriter writer; + WriteMessage(&writer, "Message"); + WriteTag(&writer, "Key1", "Value1"); + WriteTag(&writer, "Key2", "Value2"); + + TLogEvent event{ + .Category = Logger().GetCategory(), + .Level = ELogLevel::Debug, + .Family = ELogFamily::PlainText, + .Payload = writer.Finish(), + .ThreadId = 42, + .FiberId = 31, + .TraceId = TTraceId(1, 2, 3, 4), + }; + + auto formatter = std::make_unique<TPlainTextLogFormatter>(); + TStringStream stringStream; + formatter->WriteFormatted(&stringStream, event); + + // The tab-separated fields are: <date> <level> <category> <message> <thread> <fiber> <trace>. + EXPECT_EQ( + stringStream.Str().substr(DateLength), + Format("\tD\t%v\tMessage (Key1: Value1, Key2: Value2)\t2a\t1f\t4-3-2-1\n", Logger().GetCategory()->Name)); +} + +TEST_F(TLoggingTest, PlainTextFormatterWellKnownTag) +{ + TTaggedPayloadWriter writer; + WriteMessage(&writer, "Message"); + WriteTag(&writer, "Key", "Value"); + // A well-known tag (e.g. an error) with a newline in its value. + WriteWellKnownTag(&writer, "Error", "a\nb"); + + TLogEvent event{ + .Category = Logger().GetCategory(), + .Level = ELogLevel::Debug, + .Family = ELogFamily::PlainText, + .Payload = writer.Finish(), + .ThreadId = 42, + .FiberId = 31, + .TraceId = TTraceId(1, 2, 3, 4), + }; + + auto formatter = std::make_unique<TPlainTextLogFormatter>(); + TStringStream stringStream; + formatter->WriteFormatted(&stringStream, event); + + // The regular tag renders inline; the well-known tag is appended after the |(...)| + // group, separated by an escaped newline; its value's own newline is escaped too, + // so the record stays a single physical line. + EXPECT_EQ( + stringStream.Str().substr(DateLength), + Format("\tD\t%v\tMessage (Key: Value)\\na\\nb\t2a\t1f\t4-3-2-1\n", Logger().GetCategory()->Name)); +} + TEST_F(TLoggingTest, WithMinLevel) { TTempFile debugFile(GenerateLogFileName()); @@ -1551,18 +1625,18 @@ TEST_P(TLoggingTagsTest, All) if (hasMessageTag) { EXPECT_EQ( expected, - ToString(NLogging::NDetail::BuildLogMessage( + ToString(GetMessageFromTaggedPayload(NLogging::BuildLogMessage( loggingContext, logger, "Log message (Value: %v)", - 123).MessageRef)); + 123).Payload))); } else { EXPECT_EQ( expected, - ToString(NLogging::NDetail::BuildLogMessage( + ToString(GetMessageFromTaggedPayload(NLogging::BuildLogMessage( loggingContext, logger, - "Log message").MessageRef)); + "Log message").Payload))); } } @@ -1654,7 +1728,7 @@ TEST_F(TLongMessagesTest, WithoutPerThreadCache) TTempFile logFile(GenerateLogFileName()); ConfigureForLongMessages(logFile.Name()); std::thread thread([&] { - NLogging::NDetail::TMessageStringBuilder::DisablePerThreadCache(); + TTaggedPayloadWriter::DisablePerThreadCache(); LogLongMessages(); }); thread.join(); @@ -1665,8 +1739,8 @@ TEST_F(TLoggingTest, Anchors) { NLogging::TLogger logger; NLogging::TLoggingContext context{}; - EXPECT_EQ(NLogging::NDetail::BuildLogMessage(context, logger, "Simple message").Anchor, "Simple message"); - EXPECT_EQ(NLogging::NDetail::BuildLogMessage(context, logger, "Simple message (Param: %v)", 1).Anchor, "Simple message (Param: %v)"); + EXPECT_EQ(NLogging::BuildLogMessage(context, logger, "Simple message").Anchor, "Simple message"); + EXPECT_EQ(NLogging::BuildLogMessage(context, logger, "Simple message (Param: %v)", 1).Anchor, "Simple message (Param: %v)"); } //////////////////////////////////////////////////////////////////////////////// @@ -1703,7 +1777,7 @@ public: void Write(const TLogEvent& event) override { if (event.Category == Logger().GetCategory()) { - Messages_.push_back(std::string(Config_->Padding, ' ') + std::string(event.MessageRef.ToStringBuf())); + Messages_.push_back(std::string(Config_->Padding, ' ') + std::string(GetMessageFromTaggedPayload(std::get<TTaggedLogEventPayload>(event.Payload)))); } } diff --git a/yt/yt/core/ytree/serialize-inl.h b/yt/yt/core/ytree/serialize-inl.h index d73c02fe37e..d71395b34c1 100644 --- a/yt/yt/core/ytree/serialize-inl.h +++ b/yt/yt/core/ytree/serialize-inl.h @@ -8,6 +8,7 @@ #include <yt/yt/core/misc/error.h> #include <yt/yt/core/misc/collection_helpers.h> +#include <yt/yt/core/misc/protobuf_helpers.h> #include <yt/yt/core/yson/stream.h> #include <yt/yt/core/yson/string.h> @@ -517,7 +518,7 @@ void Serialize( const T& message, NYson::IYsonConsumer* consumer) { - consumer->OnStringScalar(message.SerializeAsStringOrThrow()); + consumer->OnStringScalar(SerializeProtoToString(message)); } template <class T, class TTag, TStrongTypedefOptions Options> @@ -746,7 +747,7 @@ void Deserialize( { std::string string; Deserialize(string, node); - message.ParseFromStringOrThrow(string); + DeserializeProto(&message, TRef::FromString(string)); } template <class T, class TTag, TStrongTypedefOptions Options> diff --git a/yt/yt_proto/yt/client/node_tracker_client/proto/node.proto b/yt/yt_proto/yt/client/node_tracker_client/proto/node.proto index 42ba562054c..1058837e54a 100644 --- a/yt/yt_proto/yt/client/node_tracker_client/proto/node.proto +++ b/yt/yt_proto/yt/client/node_tracker_client/proto/node.proto @@ -97,6 +97,7 @@ message TNetworkStatistics { required string network = 1; required bool throttling_reads = 2 [default = false]; + optional bool throttling_writes = 3 [default = false]; } // Describes the current state of the whole node. |