chalkruntime 3.43.74__tar.gz → 3.43.76__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/PKG-INFO +1 -1
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/invoker/vectorized_hasmany_sampler.py +19 -0
- chalkruntime-3.43.76/chalkruntime/sql_rewriter/filter_injection_site.py +167 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime.egg-info/PKG-INFO +1 -1
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/pyproject.toml +1 -1
- chalkruntime-3.43.74/chalkruntime/sql_rewriter/filter_injection_site.py +0 -382
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/README.md +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/__init__.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/constants.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/dataframe/__init__.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/dataframe/dataframe.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/dataframe/lazyframe.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/exc/__init__.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/exc/failed_argument.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/exc/resolver_errors.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/exc/wrapped_resolver_exception.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/graph/__init__.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/graph/_sql_source_pickle_patch.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/graph/chalk_overload.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/graph/convert_chalkpy_underscore.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/graph/dataframe_resolver_correlation.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/graph/feature.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/graph/feature_alias.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/graph/filter_conversion.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/graph/global_graph.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/graph/graph.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/graph/graph_impl.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/graph/graph_proxy.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/graph/graph_state.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/graph/jinja_parser.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/graph/mat_agg_filter_expression.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/graph/materializations.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/graph/materialized_feature_view.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/graph/maybe_named_collection.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/graph/named_query.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/graph/nearest_neighbor.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/graph/overlay_graph.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/graph/protograph_deserializer.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/graph/protograph_serializer.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/graph/pseudofeatures.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/graph/resolver.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/graph/singletons.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/graph/sklearn_model_parser.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/graph/source_secrets.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/graph/stream_resolver.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/graph/underscore.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/graph/underscore_codec_info.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/graph/underscore_operation_registry.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/graph/variables.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/heaptrack_launcher.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/incrementalization/__init__.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/incrementalization/group_incrementalizer.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/incrementalization/incrementalizer.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/invoker/__init__.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/invoker/batch_result_collector.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/invoker/bound_invoker.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/invoker/bound_invoker_cache.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/invoker/general_bound_invoker.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/invoker/no_arg_scalar_invoker.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/invoker/one_to_one_invoker.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/invoker/overlay_features.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/invoker/parse_external_resolver.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/invoker/partition_batch.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/invoker/query_execution_parameters.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/invoker/resolver_args_builder.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/invoker/resolver_input.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/invoker/resolver_input_upload.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/invoker/resolver_output_metadata.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/invoker/resolver_raw_output_parsing.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/invoker/resolver_result.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/invoker/resolver_runner.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/invoker/sample.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/invoker/validator.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/loader/__init__.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/loader/converter.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/loader/importer.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/memray_launcher.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/metadata.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/py.typed +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/server/__init__.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/server/branch_model_endpoint_config.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/server/config.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/server/entrypoint.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/server/env_helper.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/server/remote_python_function_registry_client.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/server/server_context_snapshot.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/server/service.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/server/subprocess_hook_service.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/sql_rewriter/__init__.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/sql_rewriter/composed_rewriter.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/sql_rewriter/contextual_query_rewriter.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/sql_rewriter/filter_query_rewriter.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/sql_rewriter/identity_rewriter.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/sql_rewriter/query_rewriter.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/sql_rewriter/query_rewriter_helper.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/sql_rewriter/sqlglot_import_lock.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/sql_rewriter/sqlglot_parsing.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/sqlalchemy_mysql_keepalive.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/streaming/__init__.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/streaming/converter_utils.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/streaming/exc.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/streaming/message_parsing.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/streaming/resolver_utils.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/streaming/streaming_invoker_kernel/__init__.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/streaming/streaming_invoker_kernel/local_streaming_invoker_kernel.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/streaming/types.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/streaming/window_keys.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/utils/__init__.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/utils/async_helpers.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/utils/contextvars.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/utils/datadog.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/utils/internal_pl_utils.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/utils/posix_file_lock.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/utils/proto_bytes_cache.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/utils/pyarrow_utils.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/utils/tracing.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/utils/viztracer_profiling.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/valgrind_launcher.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime.egg-info/SOURCES.txt +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime.egg-info/dependency_links.txt +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime.egg-info/requires.txt +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime.egg-info/top_level.txt +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/setup.cfg +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/setup.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/tests/test_branch_model_endpoint_config.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/tests/test_captured_global_function_refs.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/tests/test_datadog.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/tests/test_feature_aliases.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/tests/test_graph_cache.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/tests/test_internal_pl_utils.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/tests/test_mat_agg_filter_expressions.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/tests/test_mat_agg_projection_expressions.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/tests/test_materialized_feature_view.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/tests/test_posix_file_lock.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/tests/test_pyarrow_utils.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/tests/test_server_context_snapshot.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/tests/test_sql_max_row_version_lookback_deserialization.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/tests/test_sql_retry_policy_deserialization.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/tests/test_sqlalchemy_mysql_keepalive.py +0 -0
- {chalkruntime-3.43.74 → chalkruntime-3.43.76}/tests/test_subprocess_hook_service.py +0 -0
{chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/invoker/vectorized_hasmany_sampler.py
RENAMED
|
@@ -689,6 +689,12 @@ class PolarsVectorizedHasManySampler(VectorizedHasManySampler):
|
|
|
689
689
|
hm_lf = hm_df
|
|
690
690
|
hm_mapping_lf = mapping_table
|
|
691
691
|
|
|
692
|
+
# See `HAS_MANY_INPUT_TS_COL_NAME`: present iff each has-many row is a version for one input time,
|
|
693
|
+
# in which case an input row only matches the rows computed for its own `__ts__`.
|
|
694
|
+
join_on_input_time = HAS_MANY_INPUT_TS_COL_NAME in hm_lf.columns
|
|
695
|
+
if join_on_input_time and TS_COL_NAME not in resolver_inputs_lf.columns:
|
|
696
|
+
raise ValueError(f"The resolver inputs must have a '{TS_COL_NAME}' column to match has-many data on it")
|
|
697
|
+
|
|
692
698
|
# If we do not have the feature time column, append it with nulls
|
|
693
699
|
if foreign_ts_feature.fqn not in hm_lf.columns:
|
|
694
700
|
hm_lf = hm_lf.with_columns(pl.lit(None, pl.Datetime("us", "UTC")).alias(foreign_ts_feature.fqn))
|
|
@@ -703,6 +709,9 @@ class PolarsVectorizedHasManySampler(VectorizedHasManySampler):
|
|
|
703
709
|
unique_on = [foreign_pkey_feature.fqn, foreign_ts_feature.fqn]
|
|
704
710
|
if not allow_planner_postponed_has_many_sampling_planner_option:
|
|
705
711
|
unique_on.append(TS_COL_NAME)
|
|
712
|
+
if join_on_input_time:
|
|
713
|
+
# Versions computed for different input times are distinct rows, not duplicates.
|
|
714
|
+
unique_on.append(HAS_MANY_INPUT_TS_COL_NAME)
|
|
706
715
|
# Add __index__ to prevent cross-transaction deduplication in bulk queries
|
|
707
716
|
# This ensures that duplicate GL codes across different transactions in a bulk query
|
|
708
717
|
# are treated as distinct results, not as duplicates to be removed
|
|
@@ -726,6 +735,11 @@ class PolarsVectorizedHasManySampler(VectorizedHasManySampler):
|
|
|
726
735
|
over = [pl.col(foreign_pkey_feature.fqn)]
|
|
727
736
|
if INDEX_COL_NAME in hm_lf.columns:
|
|
728
737
|
over.append(pl.col(INDEX_COL_NAME))
|
|
738
|
+
if join_on_input_time:
|
|
739
|
+
# A row is replaced only by a later version computed for the same input time; the versions for
|
|
740
|
+
# other input times are separate histories and would otherwise mark this one as replaced at its
|
|
741
|
+
# own feature time.
|
|
742
|
+
over.append(pl.col(HAS_MANY_INPUT_TS_COL_NAME))
|
|
729
743
|
hm_lf = hm_lf.with_columns(
|
|
730
744
|
[
|
|
731
745
|
pl.col(foreign_ts_feature.fqn).alias(OBSERVED_AT_FEATURE_COL_NAME),
|
|
@@ -788,6 +802,9 @@ class PolarsVectorizedHasManySampler(VectorizedHasManySampler):
|
|
|
788
802
|
for f in filters_from_feature
|
|
789
803
|
]
|
|
790
804
|
|
|
805
|
+
if join_on_input_time:
|
|
806
|
+
left_join_cols = [*left_join_cols, TS_COL_NAME]
|
|
807
|
+
right_join_cols = [*right_join_cols, HAS_MANY_INPUT_TS_COL_NAME]
|
|
791
808
|
# Rename the left join column here since all the left table's column names had that prefix added.
|
|
792
809
|
left_join_cols = [RESOLVER_INPUT_PREFIX + left_join_col for left_join_col in left_join_cols]
|
|
793
810
|
|
|
@@ -1472,6 +1489,8 @@ class PyArrowVectorizedHasManySampler(VectorizedHasManySampler):
|
|
|
1472
1489
|
table_has_struct_or_list(resolver_inputs_table)
|
|
1473
1490
|
or table_has_struct_or_list(has_many_table[0])
|
|
1474
1491
|
or has_many_feature.underlying.get_nearest_neighbor_join_info()
|
|
1492
|
+
# This sampler does not match per-input-time has-many rows on the input time.
|
|
1493
|
+
or HAS_MANY_INPUT_TS_COL_NAME in has_many_table[0].column_names
|
|
1475
1494
|
)
|
|
1476
1495
|
|
|
1477
1496
|
|
|
@@ -0,0 +1,167 @@
|
|
|
1
|
+
"""Support for the ``chalk_filter_injection_site`` marker comment.
|
|
2
|
+
|
|
3
|
+
By default, when Chalk pushes a filter into a ``.chalk.sql`` resolver, it wraps the entire query
|
|
4
|
+
in ``SELECT * FROM (<original query>) tmp WHERE <filter>``. That outer ``WHERE`` is correct, but
|
|
5
|
+
some engines (notably Redshift) will not always push such an outer predicate down through CTEs,
|
|
6
|
+
joins, aggregations, or window functions to the underlying table scan. A resolver author can place
|
|
7
|
+
a marker comment to tell us exactly where the filter should be injected instead::
|
|
8
|
+
|
|
9
|
+
WITH foo AS (SELECT * FROM bar /* chalk_filter_injection_site */)
|
|
10
|
+
SELECT * FROM foo
|
|
11
|
+
|
|
12
|
+
The marker may optionally name the (output) columns that should be filtered at that site. This is
|
|
13
|
+
how filters get routed when several CTEs each need a different filter::
|
|
14
|
+
|
|
15
|
+
WITH
|
|
16
|
+
users AS (SELECT * FROM users_t /* chalk_filter_injection_site[user_id] */),
|
|
17
|
+
accounts AS (SELECT * FROM accounts_t /* chalk_filter_injection_site[account_id] */)
|
|
18
|
+
SELECT u.user_id, a.account_id, ...
|
|
19
|
+
FROM users u JOIN accounts a ON u.account_id = a.account_id
|
|
20
|
+
|
|
21
|
+
The named columns are the resolver's *output* column names (the keys of the resolver field
|
|
22
|
+
mapping), and they must be referenceable from the marked ``SELECT``'s ``WHERE`` clause. The common
|
|
23
|
+
``SELECT *`` pass-through case satisfies this; a renaming projection (``SELECT x AS user_id``) does
|
|
24
|
+
not, because standard SQL evaluates ``WHERE`` before ``SELECT`` aliases.
|
|
25
|
+
|
|
26
|
+
The marker only changes *where* a filter lands. The outer wrap that each rewriter already produces
|
|
27
|
+
is preserved as a correctness backstop, so a mis-placed or unsatisfiable marker never changes query
|
|
28
|
+
results -- at worst the in-place injection is skipped and behavior falls back to the outer wrap.
|
|
29
|
+
|
|
30
|
+
One placement caveat: the sqlglot *parser* drops comments in some positions (notably an own-line
|
|
31
|
+
marker between the WHERE predicate and QUALIFY), so the AST-based site finding above cannot see
|
|
32
|
+
them and this module's python rewriters keep the outer wrap. The native execution path locates the
|
|
33
|
+
marker textually instead (libchalk's ``FilterInjection.for_sql_resolver``), so it honors those placements.
|
|
34
|
+
"""
|
|
35
|
+
|
|
36
|
+
from __future__ import annotations
|
|
37
|
+
|
|
38
|
+
import dataclasses
|
|
39
|
+
import re
|
|
40
|
+
from typing import Iterable, Optional
|
|
41
|
+
|
|
42
|
+
import sqlglot
|
|
43
|
+
import sqlglot.expressions as sqlglot_expressions
|
|
44
|
+
from chalk.utils.log_with_context import get_logger
|
|
45
|
+
|
|
46
|
+
_logger = get_logger(__name__)
|
|
47
|
+
|
|
48
|
+
INJECTION_SITE_MARKER = "chalk_filter_injection_site"
|
|
49
|
+
|
|
50
|
+
# The trailing ``\b`` keeps ``chalk_filter_injection_site_other`` from matching the marker. The
|
|
51
|
+
# optional bracket group captures a comma-separated list of output column names. We ``search``
|
|
52
|
+
# (rather than fullmatch) so the marker is recognized even when it shares a comment with other text.
|
|
53
|
+
_MARKER_RE = re.compile(INJECTION_SITE_MARKER + r"\b\s*(?:\[(?P<cols>[^\]]*)\])?", re.IGNORECASE)
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
@dataclasses.dataclass(frozen=True)
|
|
57
|
+
class FilterInjectionSite:
|
|
58
|
+
"""A ``SELECT`` carrying a ``chalk_filter_injection_site`` comment, plus the optional set of
|
|
59
|
+
output columns the author scoped that site to."""
|
|
60
|
+
|
|
61
|
+
select: sqlglot_expressions.Select
|
|
62
|
+
columns: Optional[frozenset[str]]
|
|
63
|
+
"""Output columns this site accepts. ``None`` means the site is unqualified and accepts every
|
|
64
|
+
filterable column; an empty set means the site was explicitly scoped to no columns."""
|
|
65
|
+
|
|
66
|
+
def accepts(self, column_name: str) -> bool:
|
|
67
|
+
return self.columns is None or column_name in self.columns
|
|
68
|
+
|
|
69
|
+
def inject(self, condition: sqlglot_expressions.Condition | str, *, dialect: str | None) -> None:
|
|
70
|
+
"""AND ``condition`` into this site's ``WHERE`` clause, mutating the parsed tree in place.
|
|
71
|
+
|
|
72
|
+
``copy=False`` is required: the surrounding query (which the rewriter subsequently wraps)
|
|
73
|
+
holds a reference to this same ``Select`` node, so the injection must mutate it rather than
|
|
74
|
+
return a fresh copy.
|
|
75
|
+
"""
|
|
76
|
+
self.select.where(condition, dialect=dialect, copy=False)
|
|
77
|
+
|
|
78
|
+
|
|
79
|
+
def _parse_marker(comment: str) -> tuple[bool, Optional[frozenset[str]]]:
|
|
80
|
+
"""Return ``(is_marker, columns)`` for a single comment string.
|
|
81
|
+
|
|
82
|
+
``columns`` is ``None`` for an unqualified marker and a (possibly empty) frozenset of output
|
|
83
|
+
column names for a ``[...]``-qualified marker.
|
|
84
|
+
"""
|
|
85
|
+
match = _MARKER_RE.search(comment)
|
|
86
|
+
if match is None:
|
|
87
|
+
return False, None
|
|
88
|
+
cols_group = match.group("cols")
|
|
89
|
+
if cols_group is None:
|
|
90
|
+
return True, None
|
|
91
|
+
columns = frozenset(stripped for raw in cols_group.split(",") if (stripped := raw.strip().strip("\"'")))
|
|
92
|
+
return True, columns
|
|
93
|
+
|
|
94
|
+
|
|
95
|
+
def _iter_nodes(expression: sqlglot.Expression) -> Iterable[sqlglot.Expression]:
|
|
96
|
+
# ``Expression.walk`` yields ``(node, parent, key)`` tuples on the pinned sqlglot (19.8.3) but
|
|
97
|
+
# yields bare nodes on newer releases; tolerate both so this keeps working across upgrades.
|
|
98
|
+
for item in expression.walk():
|
|
99
|
+
yield item[0] if isinstance(item, tuple) else item
|
|
100
|
+
|
|
101
|
+
|
|
102
|
+
def find_filter_injection_sites(expression: sqlglot.Expression) -> list[FilterInjectionSite]:
|
|
103
|
+
"""Find every ``chalk_filter_injection_site`` marker in ``expression``.
|
|
104
|
+
|
|
105
|
+
Multiple markers that resolve to the same ``SELECT`` are merged into a single site (their column
|
|
106
|
+
scopes are unioned; an unqualified marker widens the merged site to accept all columns) so a
|
|
107
|
+
given condition is never injected into the same ``SELECT`` twice.
|
|
108
|
+
"""
|
|
109
|
+
by_select: dict[int, FilterInjectionSite] = {}
|
|
110
|
+
order: list[int] = []
|
|
111
|
+
for node in _iter_nodes(expression):
|
|
112
|
+
comments = node.comments
|
|
113
|
+
if not comments:
|
|
114
|
+
continue
|
|
115
|
+
for comment in comments:
|
|
116
|
+
is_marker, columns = _parse_marker(comment)
|
|
117
|
+
if not is_marker:
|
|
118
|
+
continue
|
|
119
|
+
select = (
|
|
120
|
+
node if isinstance(node, sqlglot_expressions.Select) else node.find_ancestor(sqlglot_expressions.Select)
|
|
121
|
+
)
|
|
122
|
+
if select is None:
|
|
123
|
+
_logger.warning(
|
|
124
|
+
"Ignoring a '%s' marker that is not contained within a SELECT statement.",
|
|
125
|
+
INJECTION_SITE_MARKER,
|
|
126
|
+
)
|
|
127
|
+
continue
|
|
128
|
+
key = id(select)
|
|
129
|
+
if key not in by_select:
|
|
130
|
+
by_select[key] = FilterInjectionSite(select=select, columns=columns)
|
|
131
|
+
order.append(key)
|
|
132
|
+
else:
|
|
133
|
+
existing = by_select[key]
|
|
134
|
+
if existing.columns is None or columns is None:
|
|
135
|
+
merged: Optional[frozenset[str]] = None
|
|
136
|
+
else:
|
|
137
|
+
merged = existing.columns | columns
|
|
138
|
+
by_select[key] = FilterInjectionSite(select=select, columns=merged)
|
|
139
|
+
return [by_select[key] for key in order]
|
|
140
|
+
|
|
141
|
+
|
|
142
|
+
def inject_at_accepting_sites(
|
|
143
|
+
sites: list[FilterInjectionSite],
|
|
144
|
+
column_name: str,
|
|
145
|
+
condition: sqlglot_expressions.Condition | str,
|
|
146
|
+
*,
|
|
147
|
+
dialect: str | None,
|
|
148
|
+
) -> int:
|
|
149
|
+
"""Inject ``condition`` into every site that accepts ``column_name``; return the count injected.
|
|
150
|
+
|
|
151
|
+
A sqlglot condition is copied for each site so the same node is never shared across two places in
|
|
152
|
+
the tree (raw ``WHERE`` strings are re-parsed by ``Select.where`` and so need no copy).
|
|
153
|
+
"""
|
|
154
|
+
injected = 0
|
|
155
|
+
for site in sites:
|
|
156
|
+
if not site.accepts(column_name):
|
|
157
|
+
continue
|
|
158
|
+
to_inject = condition.copy() if isinstance(condition, sqlglot_expressions.Expression) else condition
|
|
159
|
+
site.inject(to_inject, dialect=dialect)
|
|
160
|
+
injected += 1
|
|
161
|
+
return injected
|
|
162
|
+
|
|
163
|
+
|
|
164
|
+
def is_noop_condition(condition: sqlglot.Expression) -> bool:
|
|
165
|
+
"""Whether ``condition`` is a literal ``TRUE`` (what the filter converters emit when nothing can
|
|
166
|
+
be pushed down). Injecting it would only add ``WHERE TRUE`` noise, so callers skip it."""
|
|
167
|
+
return isinstance(condition, sqlglot_expressions.Boolean) and condition.this is True
|
|
@@ -1,382 +0,0 @@
|
|
|
1
|
-
"""Support for the ``chalk_filter_injection_site`` marker comment.
|
|
2
|
-
|
|
3
|
-
By default, when Chalk pushes a filter into a ``.chalk.sql`` resolver, it wraps the entire query
|
|
4
|
-
in ``SELECT * FROM (<original query>) tmp WHERE <filter>``. That outer ``WHERE`` is correct, but
|
|
5
|
-
some engines (notably Redshift) will not always push such an outer predicate down through CTEs,
|
|
6
|
-
joins, aggregations, or window functions to the underlying table scan. A resolver author can place
|
|
7
|
-
a marker comment to tell us exactly where the filter should be injected instead::
|
|
8
|
-
|
|
9
|
-
WITH foo AS (SELECT * FROM bar /* chalk_filter_injection_site */)
|
|
10
|
-
SELECT * FROM foo
|
|
11
|
-
|
|
12
|
-
The marker may optionally name the (output) columns that should be filtered at that site. This is
|
|
13
|
-
how filters get routed when several CTEs each need a different filter::
|
|
14
|
-
|
|
15
|
-
WITH
|
|
16
|
-
users AS (SELECT * FROM users_t /* chalk_filter_injection_site[user_id] */),
|
|
17
|
-
accounts AS (SELECT * FROM accounts_t /* chalk_filter_injection_site[account_id] */)
|
|
18
|
-
SELECT u.user_id, a.account_id, ...
|
|
19
|
-
FROM users u JOIN accounts a ON u.account_id = a.account_id
|
|
20
|
-
|
|
21
|
-
The named columns are the resolver's *output* column names (the keys of the resolver field
|
|
22
|
-
mapping), and they must be referenceable from the marked ``SELECT``'s ``WHERE`` clause. The common
|
|
23
|
-
``SELECT *`` pass-through case satisfies this; a renaming projection (``SELECT x AS user_id``) does
|
|
24
|
-
not, because standard SQL evaluates ``WHERE`` before ``SELECT`` aliases.
|
|
25
|
-
|
|
26
|
-
The marker only changes *where* a filter lands. The outer wrap that each rewriter already produces
|
|
27
|
-
is preserved as a correctness backstop, so a mis-placed or unsatisfiable marker never changes query
|
|
28
|
-
results -- at worst the in-place injection is skipped and behavior falls back to the outer wrap.
|
|
29
|
-
|
|
30
|
-
One placement caveat: the sqlglot *parser* drops comments in some positions (notably an own-line
|
|
31
|
-
marker between the WHERE predicate and QUALIFY), so the AST-based site finding above cannot see
|
|
32
|
-
them and this module's python rewriters keep the outer wrap. ``compute_native_injection_site``
|
|
33
|
-
below locates the marker textually instead, so the native execution path honors those placements.
|
|
34
|
-
"""
|
|
35
|
-
|
|
36
|
-
from __future__ import annotations
|
|
37
|
-
|
|
38
|
-
import dataclasses
|
|
39
|
-
import re
|
|
40
|
-
from typing import Iterable, Optional
|
|
41
|
-
|
|
42
|
-
import sqlglot
|
|
43
|
-
import sqlglot.expressions as sqlglot_expressions
|
|
44
|
-
from chalk.utils.log_with_context import get_logger
|
|
45
|
-
|
|
46
|
-
from chalkruntime.sql_rewriter.sqlglot_parsing import sqlglot_dialect_object
|
|
47
|
-
|
|
48
|
-
_logger = get_logger(__name__)
|
|
49
|
-
|
|
50
|
-
INJECTION_SITE_MARKER = "chalk_filter_injection_site"
|
|
51
|
-
|
|
52
|
-
# The trailing ``\b`` keeps ``chalk_filter_injection_site_other`` from matching the marker. The
|
|
53
|
-
# optional bracket group captures a comma-separated list of output column names. We ``search``
|
|
54
|
-
# (rather than fullmatch) so the marker is recognized even when it shares a comment with other text.
|
|
55
|
-
_MARKER_RE = re.compile(INJECTION_SITE_MARKER + r"\b\s*(?:\[(?P<cols>[^\]]*)\])?", re.IGNORECASE)
|
|
56
|
-
|
|
57
|
-
|
|
58
|
-
@dataclasses.dataclass(frozen=True)
|
|
59
|
-
class FilterInjectionSite:
|
|
60
|
-
"""A ``SELECT`` carrying a ``chalk_filter_injection_site`` comment, plus the optional set of
|
|
61
|
-
output columns the author scoped that site to."""
|
|
62
|
-
|
|
63
|
-
select: sqlglot_expressions.Select
|
|
64
|
-
columns: Optional[frozenset[str]]
|
|
65
|
-
"""Output columns this site accepts. ``None`` means the site is unqualified and accepts every
|
|
66
|
-
filterable column; an empty set means the site was explicitly scoped to no columns."""
|
|
67
|
-
|
|
68
|
-
def accepts(self, column_name: str) -> bool:
|
|
69
|
-
return self.columns is None or column_name in self.columns
|
|
70
|
-
|
|
71
|
-
def inject(self, condition: sqlglot_expressions.Condition | str, *, dialect: str | None) -> None:
|
|
72
|
-
"""AND ``condition`` into this site's ``WHERE`` clause, mutating the parsed tree in place.
|
|
73
|
-
|
|
74
|
-
``copy=False`` is required: the surrounding query (which the rewriter subsequently wraps)
|
|
75
|
-
holds a reference to this same ``Select`` node, so the injection must mutate it rather than
|
|
76
|
-
return a fresh copy.
|
|
77
|
-
"""
|
|
78
|
-
self.select.where(condition, dialect=dialect, copy=False)
|
|
79
|
-
|
|
80
|
-
|
|
81
|
-
def _parse_marker(comment: str) -> tuple[bool, Optional[frozenset[str]]]:
|
|
82
|
-
"""Return ``(is_marker, columns)`` for a single comment string.
|
|
83
|
-
|
|
84
|
-
``columns`` is ``None`` for an unqualified marker and a (possibly empty) frozenset of output
|
|
85
|
-
column names for a ``[...]``-qualified marker.
|
|
86
|
-
"""
|
|
87
|
-
match = _MARKER_RE.search(comment)
|
|
88
|
-
if match is None:
|
|
89
|
-
return False, None
|
|
90
|
-
cols_group = match.group("cols")
|
|
91
|
-
if cols_group is None:
|
|
92
|
-
return True, None
|
|
93
|
-
columns = frozenset(stripped for raw in cols_group.split(",") if (stripped := raw.strip().strip("\"'")))
|
|
94
|
-
return True, columns
|
|
95
|
-
|
|
96
|
-
|
|
97
|
-
def _iter_nodes(expression: sqlglot.Expression) -> Iterable[sqlglot.Expression]:
|
|
98
|
-
# ``Expression.walk`` yields ``(node, parent, key)`` tuples on the pinned sqlglot (19.8.3) but
|
|
99
|
-
# yields bare nodes on newer releases; tolerate both so this keeps working across upgrades.
|
|
100
|
-
for item in expression.walk():
|
|
101
|
-
yield item[0] if isinstance(item, tuple) else item
|
|
102
|
-
|
|
103
|
-
|
|
104
|
-
def find_filter_injection_sites(expression: sqlglot.Expression) -> list[FilterInjectionSite]:
|
|
105
|
-
"""Find every ``chalk_filter_injection_site`` marker in ``expression``.
|
|
106
|
-
|
|
107
|
-
Multiple markers that resolve to the same ``SELECT`` are merged into a single site (their column
|
|
108
|
-
scopes are unioned; an unqualified marker widens the merged site to accept all columns) so a
|
|
109
|
-
given condition is never injected into the same ``SELECT`` twice.
|
|
110
|
-
"""
|
|
111
|
-
by_select: dict[int, FilterInjectionSite] = {}
|
|
112
|
-
order: list[int] = []
|
|
113
|
-
for node in _iter_nodes(expression):
|
|
114
|
-
comments = node.comments
|
|
115
|
-
if not comments:
|
|
116
|
-
continue
|
|
117
|
-
for comment in comments:
|
|
118
|
-
is_marker, columns = _parse_marker(comment)
|
|
119
|
-
if not is_marker:
|
|
120
|
-
continue
|
|
121
|
-
select = (
|
|
122
|
-
node if isinstance(node, sqlglot_expressions.Select) else node.find_ancestor(sqlglot_expressions.Select)
|
|
123
|
-
)
|
|
124
|
-
if select is None:
|
|
125
|
-
_logger.warning(
|
|
126
|
-
"Ignoring a '%s' marker that is not contained within a SELECT statement.",
|
|
127
|
-
INJECTION_SITE_MARKER,
|
|
128
|
-
)
|
|
129
|
-
continue
|
|
130
|
-
key = id(select)
|
|
131
|
-
if key not in by_select:
|
|
132
|
-
by_select[key] = FilterInjectionSite(select=select, columns=columns)
|
|
133
|
-
order.append(key)
|
|
134
|
-
else:
|
|
135
|
-
existing = by_select[key]
|
|
136
|
-
if existing.columns is None or columns is None:
|
|
137
|
-
merged: Optional[frozenset[str]] = None
|
|
138
|
-
else:
|
|
139
|
-
merged = existing.columns | columns
|
|
140
|
-
by_select[key] = FilterInjectionSite(select=select, columns=merged)
|
|
141
|
-
return [by_select[key] for key in order]
|
|
142
|
-
|
|
143
|
-
|
|
144
|
-
def inject_at_accepting_sites(
|
|
145
|
-
sites: list[FilterInjectionSite],
|
|
146
|
-
column_name: str,
|
|
147
|
-
condition: sqlglot_expressions.Condition | str,
|
|
148
|
-
*,
|
|
149
|
-
dialect: str | None,
|
|
150
|
-
) -> int:
|
|
151
|
-
"""Inject ``condition`` into every site that accepts ``column_name``; return the count injected.
|
|
152
|
-
|
|
153
|
-
A sqlglot condition is copied for each site so the same node is never shared across two places in
|
|
154
|
-
the tree (raw ``WHERE`` strings are re-parsed by ``Select.where`` and so need no copy).
|
|
155
|
-
"""
|
|
156
|
-
injected = 0
|
|
157
|
-
for site in sites:
|
|
158
|
-
if not site.accepts(column_name):
|
|
159
|
-
continue
|
|
160
|
-
to_inject = condition.copy() if isinstance(condition, sqlglot_expressions.Expression) else condition
|
|
161
|
-
site.inject(to_inject, dialect=dialect)
|
|
162
|
-
injected += 1
|
|
163
|
-
return injected
|
|
164
|
-
|
|
165
|
-
|
|
166
|
-
def is_noop_condition(condition: sqlglot.Expression) -> bool:
|
|
167
|
-
"""Whether ``condition`` is a literal ``TRUE`` (what the filter converters emit when nothing can
|
|
168
|
-
be pushed down). Injecting it would only add ``WHERE TRUE`` noise, so callers skip it."""
|
|
169
|
-
return isinstance(condition, sqlglot_expressions.Boolean) and condition.this is True
|
|
170
|
-
|
|
171
|
-
|
|
172
|
-
@dataclasses.dataclass(frozen=True)
|
|
173
|
-
class NativeFilterInjectionSite:
|
|
174
|
-
"""A pre-computed splice point for the native (libchalk) rewriters.
|
|
175
|
-
|
|
176
|
-
The native execution path keeps the resolver's SQL as verbatim text and cannot re-render it
|
|
177
|
-
through sqlglot, so instead of handing it a parsed site we hand it a *text offset*:
|
|
178
|
-
|
|
179
|
-
- ``inside_where`` False: ``index`` is where a new `` WHERE (<filter>)`` clause belongs -- after
|
|
180
|
-
the marked SELECT's FROM/JOIN zone (and after the marker comment itself), before any
|
|
181
|
-
GROUP BY/QUALIFY/ORDER BY/closing parenthesis.
|
|
182
|
-
- ``inside_where`` True: ``index`` points immediately past the marked SELECT's existing WHERE
|
|
183
|
-
keyword, so filters are spliced as ``(<filter>) AND`` prefixes.
|
|
184
|
-
|
|
185
|
-
The offsets are computed against the exact ``sql`` string passed in, which must be the same
|
|
186
|
-
string handed to the native UDF as ``raw_query``.
|
|
187
|
-
"""
|
|
188
|
-
|
|
189
|
-
index: int
|
|
190
|
-
inside_where: bool
|
|
191
|
-
|
|
192
|
-
|
|
193
|
-
# Token types that end the zone where a WHERE clause may be introduced in the marked SELECT.
|
|
194
|
-
_CLAUSE_STOP_TOKENS = frozenset(
|
|
195
|
-
name
|
|
196
|
-
for name in (
|
|
197
|
-
"GROUP_BY",
|
|
198
|
-
"HAVING",
|
|
199
|
-
"QUALIFY",
|
|
200
|
-
"WINDOW",
|
|
201
|
-
"ORDER_BY",
|
|
202
|
-
"LIMIT",
|
|
203
|
-
"OFFSET",
|
|
204
|
-
"FETCH",
|
|
205
|
-
"UNION",
|
|
206
|
-
"EXCEPT",
|
|
207
|
-
"INTERSECT",
|
|
208
|
-
"SEMICOLON",
|
|
209
|
-
)
|
|
210
|
-
if hasattr(sqlglot.TokenType, name)
|
|
211
|
-
)
|
|
212
|
-
|
|
213
|
-
|
|
214
|
-
def _enclosing_comment_span(sql: str, name_match: re.Match[str]) -> Optional[tuple[int, int]]:
|
|
215
|
-
"""The (start, end) offsets of the comment containing the marker name at ``name_match`` --
|
|
216
|
-
through the ``*/`` of a block comment, or through the newline of a ``--`` line comment."""
|
|
217
|
-
block_start = sql.rfind("/*", 0, name_match.start())
|
|
218
|
-
line_start = sql.rfind("--", 0, name_match.start())
|
|
219
|
-
if block_start == -1 and line_start == -1:
|
|
220
|
-
return None
|
|
221
|
-
if block_start > line_start:
|
|
222
|
-
block_end = sql.find("*/", name_match.end())
|
|
223
|
-
return None if block_end == -1 else (block_start, block_end + len("*/"))
|
|
224
|
-
newline = sql.find("\n", name_match.end())
|
|
225
|
-
# A line-comment marker with no trailing newline leaves nowhere safe to splice.
|
|
226
|
-
return None if newline == -1 else (line_start, newline + 1)
|
|
227
|
-
|
|
228
|
-
|
|
229
|
-
def compute_native_injection_site(sql: str, dialect: str | None) -> Optional[NativeFilterInjectionSite]:
|
|
230
|
-
"""Compute the text splice point for ``sql``'s ``chalk_filter_injection_site`` marker, for the
|
|
231
|
-
native SQL execution path. Returns ``None`` whenever injection cannot be done safely, in which
|
|
232
|
-
case the native rewriters fall back to their outer subquery wrap:
|
|
233
|
-
|
|
234
|
-
- no marker, several markers, or a column-qualified marker (the native rewriters build one
|
|
235
|
-
combined filter over all pushdown columns, so a scoped site cannot be honored);
|
|
236
|
-
- a marker in a GROUP BY/ORDER BY/... clause of its SELECT: no WHERE clause may start there.
|
|
237
|
-
(A marker anywhere in the FROM zone or the WHERE zone is fine: FROM-zone sites splice a new
|
|
238
|
-
WHERE at the end of the zone, WHERE-zone sites splice ``((filter)) AND`` at the clause head.);
|
|
239
|
-
- an existing WHERE predicate containing a top-level OR: the splice prepends ``(filter) AND``,
|
|
240
|
-
which would bind tighter than the OR and silently change the predicate's meaning;
|
|
241
|
-
- ``${...}`` resolver parameters (the engine substitutes them before execution, shifting every
|
|
242
|
-
offset this function computes);
|
|
243
|
-
- anything the tokenizer or parser cannot make sense of.
|
|
244
|
-
"""
|
|
245
|
-
# Cheap pre-check first: this runs for every native SQL resolver at plan time, and the
|
|
246
|
-
# overwhelmingly common case is "no marker", which must not pay for a parse.
|
|
247
|
-
if INJECTION_SITE_MARKER not in sql.lower():
|
|
248
|
-
return None
|
|
249
|
-
if "${" in sql:
|
|
250
|
-
return None
|
|
251
|
-
# Callers pass the *native* dialect names, which sqlglot sometimes spells differently, and a
|
|
252
|
-
# couple of which it cannot read at all; `sqlglot_dialect_object` answers None for those and
|
|
253
|
-
# the marker falls back to the wrap.
|
|
254
|
-
dialect_obj = sqlglot_dialect_object(dialect)
|
|
255
|
-
if dialect_obj is None:
|
|
256
|
-
return None
|
|
257
|
-
try:
|
|
258
|
-
# Everything here is textual/tokenizer-based on purpose: the sqlglot *parser* drops
|
|
259
|
-
# comments in some positions (an own-line marker between the WHERE predicate and QUALIFY),
|
|
260
|
-
# and the tokenizer attaches an own-line comment to the *following* token -- so neither
|
|
261
|
-
# AST comments nor token attachment can locate the marker reliably. The regex locates it
|
|
262
|
-
# positionally; the token stream is used only for clause/parenthesis structure.
|
|
263
|
-
name_matches = list(_MARKER_RE.finditer(sql))
|
|
264
|
-
if len(name_matches) != 1:
|
|
265
|
-
return None
|
|
266
|
-
if name_matches[0].group("cols") is not None:
|
|
267
|
-
# Column-qualified sites cannot be honored natively: the rewriters build one combined
|
|
268
|
-
# filter over all pushdown columns.
|
|
269
|
-
return None
|
|
270
|
-
comment_span = _enclosing_comment_span(sql, name_matches[0])
|
|
271
|
-
if comment_span is None:
|
|
272
|
-
return None
|
|
273
|
-
marker_start, marker_end = comment_span
|
|
274
|
-
|
|
275
|
-
tokens = dialect_obj.tokenize(sql)
|
|
276
|
-
# Anchor on the last token that ends before the marker comment begins (purely positional,
|
|
277
|
-
# immune to the tokenizer's comment-attachment quirks).
|
|
278
|
-
marker_idx = None
|
|
279
|
-
for i, token in enumerate(tokens):
|
|
280
|
-
if token.end < marker_start:
|
|
281
|
-
marker_idx = i
|
|
282
|
-
else:
|
|
283
|
-
break
|
|
284
|
-
if marker_idx is None:
|
|
285
|
-
# Nothing precedes the marker; there is no SELECT for it to mark.
|
|
286
|
-
return None
|
|
287
|
-
|
|
288
|
-
zone, zone_where_idx = _marker_clause_zone(tokens, marker_idx)
|
|
289
|
-
if zone == "WHERE":
|
|
290
|
-
# A marker anywhere in the WHERE zone authorizes the site; the splice lands at the
|
|
291
|
-
# *head* of the clause (`WHERE ((filter)) AND <original predicate>`), which stays valid
|
|
292
|
-
# SQL no matter where in the predicate the marker itself sits.
|
|
293
|
-
if zone_where_idx is None:
|
|
294
|
-
# Unreachable: a "WHERE" zone always carries its WHERE token index. Guarded rather
|
|
295
|
-
# than asserted because this runs at plan time in production (-O strips asserts).
|
|
296
|
-
return None
|
|
297
|
-
if _where_zone_has_top_level_or(tokens, zone_where_idx):
|
|
298
|
-
return None
|
|
299
|
-
return NativeFilterInjectionSite(index=tokens[zone_where_idx].end + 1, inside_where=True)
|
|
300
|
-
if zone != "FROM":
|
|
301
|
-
# GROUP BY / HAVING / QUALIFY / ORDER BY / ... -- no WHERE clause may start here.
|
|
302
|
-
return None
|
|
303
|
-
|
|
304
|
-
# FROM zone: scan forward for the end of the SELECT's FROM/WHERE region, tracking
|
|
305
|
-
# parenthesis depth relative to the marker. A closing parenthesis that drops below the
|
|
306
|
-
# marker's depth is the end of the enclosing CTE/subquery.
|
|
307
|
-
depth = 0
|
|
308
|
-
previous_end = max(tokens[marker_idx].end + 1, marker_end)
|
|
309
|
-
for token_idx in range(marker_idx + 1, len(tokens)):
|
|
310
|
-
token = tokens[token_idx]
|
|
311
|
-
token_name = token.token_type.name
|
|
312
|
-
if token_name == "L_PAREN":
|
|
313
|
-
depth += 1
|
|
314
|
-
elif token_name == "R_PAREN":
|
|
315
|
-
depth -= 1
|
|
316
|
-
if depth < 0:
|
|
317
|
-
return NativeFilterInjectionSite(index=previous_end, inside_where=False)
|
|
318
|
-
elif depth == 0:
|
|
319
|
-
if token_name == "WHERE":
|
|
320
|
-
if _where_zone_has_top_level_or(tokens, token_idx):
|
|
321
|
-
return None
|
|
322
|
-
return NativeFilterInjectionSite(index=token.end + 1, inside_where=True)
|
|
323
|
-
if token_name in _CLAUSE_STOP_TOKENS:
|
|
324
|
-
return NativeFilterInjectionSite(index=previous_end, inside_where=False)
|
|
325
|
-
previous_end = max(token.end + 1, marker_end)
|
|
326
|
-
# End of statement: the FROM zone runs to the end of the text.
|
|
327
|
-
return NativeFilterInjectionSite(index=previous_end, inside_where=False)
|
|
328
|
-
except Exception:
|
|
329
|
-
_logger.warning("Failed to compute a native filter injection site; falling back to the wrap.", exc_info=True)
|
|
330
|
-
return None
|
|
331
|
-
|
|
332
|
-
|
|
333
|
-
# Clause keywords that name the zone a backward scan lands in; anything here other than WHERE
|
|
334
|
-
# means the marker sits after the point where a WHERE clause could start. (Membership-only, so
|
|
335
|
-
# ordering is irrelevant; frozenset spelled out to satisfy the no-set-literals semgrep rule.)
|
|
336
|
-
_POST_WHERE_ZONE_TOKENS = frozenset(name for name in _CLAUSE_STOP_TOKENS if name != "SEMICOLON")
|
|
337
|
-
|
|
338
|
-
|
|
339
|
-
def _marker_clause_zone(tokens: list, marker_idx: int) -> tuple[str, Optional[int]]:
|
|
340
|
-
"""Which clause of its enclosing SELECT the marker-carrying token sits in: ``"FROM"`` (anywhere
|
|
341
|
-
from the select list through the FROM/JOIN zone), ``"WHERE"`` (paired with the index of that
|
|
342
|
-
WHERE token), or ``"POST_WHERE"`` (GROUP BY and later). Scans backward at balanced parenthesis
|
|
343
|
-
depth; leaving the enclosing subquery or finding SELECT/FROM at depth 0 means no later clause
|
|
344
|
-
keyword preceded the marker."""
|
|
345
|
-
depth = 0
|
|
346
|
-
for i in range(marker_idx, -1, -1):
|
|
347
|
-
token_name = tokens[i].token_type.name
|
|
348
|
-
if token_name == "R_PAREN":
|
|
349
|
-
depth += 1
|
|
350
|
-
elif token_name == "L_PAREN":
|
|
351
|
-
depth -= 1
|
|
352
|
-
if depth < 0:
|
|
353
|
-
return "FROM", None
|
|
354
|
-
elif depth == 0:
|
|
355
|
-
if token_name == "WHERE":
|
|
356
|
-
return "WHERE", i
|
|
357
|
-
if token_name in _POST_WHERE_ZONE_TOKENS:
|
|
358
|
-
return "POST_WHERE", None
|
|
359
|
-
if token_name in ("FROM", "SELECT"):
|
|
360
|
-
return "FROM", None
|
|
361
|
-
return "FROM", None
|
|
362
|
-
|
|
363
|
-
|
|
364
|
-
def _where_zone_has_top_level_or(tokens: list, where_idx: int) -> bool:
|
|
365
|
-
"""Whether the WHERE clause starting at ``tokens[where_idx]`` contains an OR at parenthesis
|
|
366
|
-
depth 0. Splicing ``(filter) AND`` at the head of such a predicate would bind tighter than the
|
|
367
|
-
OR and silently change its meaning, so those sites are refused."""
|
|
368
|
-
depth = 0
|
|
369
|
-
for i in range(where_idx + 1, len(tokens)):
|
|
370
|
-
token_name = tokens[i].token_type.name
|
|
371
|
-
if token_name == "L_PAREN":
|
|
372
|
-
depth += 1
|
|
373
|
-
elif token_name == "R_PAREN":
|
|
374
|
-
depth -= 1
|
|
375
|
-
if depth < 0:
|
|
376
|
-
return False
|
|
377
|
-
elif depth == 0:
|
|
378
|
-
if token_name in _CLAUSE_STOP_TOKENS:
|
|
379
|
-
return False
|
|
380
|
-
if token_name == "OR":
|
|
381
|
-
return True
|
|
382
|
-
return False
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/exc/wrapped_resolver_exception.py
RENAMED
|
File without changes
|
|
File without changes
|
{chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/graph/_sql_source_pickle_patch.py
RENAMED
|
File without changes
|
|
File without changes
|
{chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/graph/convert_chalkpy_underscore.py
RENAMED
|
File without changes
|
{chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/graph/dataframe_resolver_correlation.py
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/graph/mat_agg_filter_expression.py
RENAMED
|
File without changes
|
|
File without changes
|
{chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/graph/materialized_feature_view.py
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/graph/underscore_operation_registry.py
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/incrementalization/incrementalizer.py
RENAMED
|
File without changes
|
|
File without changes
|
{chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/invoker/batch_result_collector.py
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/invoker/parse_external_resolver.py
RENAMED
|
File without changes
|
|
File without changes
|
{chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/invoker/query_execution_parameters.py
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/invoker/resolver_output_metadata.py
RENAMED
|
File without changes
|
{chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/invoker/resolver_raw_output_parsing.py
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/server/branch_model_endpoint_config.py
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/server/server_context_snapshot.py
RENAMED
|
File without changes
|
|
File without changes
|
{chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/server/subprocess_hook_service.py
RENAMED
|
File without changes
|
|
File without changes
|
{chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/sql_rewriter/composed_rewriter.py
RENAMED
|
File without changes
|
{chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/sql_rewriter/contextual_query_rewriter.py
RENAMED
|
File without changes
|
{chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/sql_rewriter/filter_query_rewriter.py
RENAMED
|
File without changes
|
{chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/sql_rewriter/identity_rewriter.py
RENAMED
|
File without changes
|
|
File without changes
|
{chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/sql_rewriter/query_rewriter_helper.py
RENAMED
|
File without changes
|
{chalkruntime-3.43.74 → chalkruntime-3.43.76}/chalkruntime/sql_rewriter/sqlglot_import_lock.py
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{chalkruntime-3.43.74 → chalkruntime-3.43.76}/tests/test_sql_retry_policy_deserialization.py
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|