codeant-ai-for-open-source[bot] commented on code in PR #43758:
URL: https://github.com/apache/superset/pull/43758#discussion_r4088965820
##########
superset/connectors/sqla/partition_mapping.py:
##########
@@ -162,6 +235,260 @@ def _find_niladic_calls(statement: SQLStatement) ->
set[str]:
return NON_DETERMINISTIC_WHEN_NILADIC & statement.get_niladic_functions()
+def resolve_partition_mapping(datasource: SqlaTable) -> PartitionMapping |
None:
+ """
+ Resolve the dataset's mapping, or ``None`` when nothing may be mirrored.
+
+ Every bail-out here is defensive as well as functional: save-time
validation
+ rejects most of these, but rows predating the validation can still violate
+ the invariants, and a column sync can invalidate a mapping that was fine
+ when it was written.
+ """
+ if not feature_flag_manager.is_feature_enabled(FEATURE_FLAG):
+ return None
+
+ partition_column = getattr(datasource, "partition_column", None)
+ if not partition_column:
+ return None
+
+ columns_by_name = {column.column_name: column for column in
datasource.columns}
+ if partition_column not in columns_by_name:
+ # The partition column was dropped by a column sync or at the source.
+ return None
+
+ mapped_column_name = (
+ getattr(datasource, "partition_mapped_column", None) or
datasource.main_dttm_col
+ )
+ if not mapped_column_name or mapped_column_name not in columns_by_name:
+ return None
+
+ if mapped_column_name == partition_column:
+ # Self-mapping: the mirrored predicate would duplicate the original.
+ return None
+
+ mapped_column = columns_by_name[mapped_column_name]
+ transform = getattr(mapped_column, "partition_value_transform", None)
+ if not _transform_is_usable(transform, datasource.database.backend):
+ return None
+
+ if _has_active_advanced_data_type(mapped_column):
+ # `translate_filter` builds its own predicate shape from *translated*
+ # values, so the `(operator, value)` pair the operator matrix reasons
+ # about does not exist and mirroring would apply the wrong values.
+ return None
+
+ return PartitionMapping(
+ partition_column=str(partition_column),
+ mapped_column=str(mapped_column_name),
+ value_transform=cast(str, transform),
+ is_monotonic=bool(
+ getattr(mapped_column, "partition_transform_is_monotonic", False)
+ ),
+ )
+
+
+def _transform_is_usable(transform: str | None, engine: str) -> bool:
+ """
+ Whether the transform is safe to evaluate and mirror through.
+
+ Mirrors the Tier-2 half of `validate_partition_mapping` plus the Jinja
+ block, so a mapping saved before a check existed -- or one whose engine
+ changed underneath it -- is still skipped at query time.
+ """
+ if not transform or not transform.strip():
+ return False
+ if not contains_value_placeholder(transform):
+ return False
+ if contains_jinja(transform):
+ return False
+ return is_parseable(transform, engine)
+
+
+def _has_active_advanced_data_type(column: TableColumn) -> bool:
+ advanced_data_type = getattr(column, "advanced_data_type", None)
+ if not advanced_data_type:
+ return False
+ if not
feature_flag_manager.is_feature_enabled("ENABLE_ADVANCED_DATA_TYPES"):
+ return False
+ return advanced_data_type in app.config.get("ADVANCED_DATA_TYPES", {})
+
+
+def build_probe_sql(
+ transform: str,
+ values: list[Any],
+ dialect: Dialect | None = None,
+) -> str:
+ """
+ Compile a single ``SELECT`` that evaluates the transform at every value.
+
+ Values are attacker-controlled (a Gamma user picks filter values), so they
+ are bound as parameters and rendered by the dialect's own literal processor
+ rather than interpolated into the SQL text.
+
+ Note this deliberately does *not* go through ``BaseEngineSpec``'s text
+ helper, which escapes ``:`` on every engine but Athena and would destroy
the
+ ``:value`` placeholder before it can be bound.
+ """
+ selections = []
+ for index, value in enumerate(values):
+ clause = sa.text(transform).bindparams(sa.bindparam("value",
value=value))
+ compiled = clause.compile(
+ dialect=dialect,
+ compile_kwargs={"literal_binds": True},
+ )
+ selections.append(f"{compiled} AS v{index}")
+ return "SELECT " + ", ".join(selections)
+
+
+def evaluate_transform(
+ database: Database,
+ catalog: str | None,
+ schema: str | None,
+ transform: str,
+ values: list[Any],
+) -> list[Any] | None:
+ """
+ Evaluate ``transform`` against the engine once per distinct value.
+
+ Returns one result per input value, positionally aligned with ``values``,
or
+ ``None`` if anything at all goes wrong. Failing open costs pruning, never
+ correctness: the chart query still runs, it just scans more partitions.
+
+ The probe is pinned to the dataset's catalog and schema so session settings
+ match the chart query as closely as the connection pool allows. It still
+ runs in a *different* session, which is why transforms calling
+ session-dependent functions are rejected at save time.
+ """
+ if not values:
+ return None
+
+ # Dedupe so a 200-value `IN` list costs one column, not 200.
+ distinct: list[Any] = []
+ seen: set[Any] = set()
+ for value in values:
+ key = _hashable(value)
+ if key not in seen:
+ seen.add(key)
+ distinct.append(value)
+
+ cache_key = _probe_cache_key(database, catalog, schema, transform,
distinct)
+ cached = _cache_get(cache_key)
+ if cached is None:
+ cached = _run_probe(database, catalog, schema, transform, distinct)
+ if cached is None:
+ # Deliberately not cached: a transient engine blip would otherwise
+ # keep the dataset pruning-free for the whole cache timeout.
+ return None
+ _cache_set(cache_key, cached)
+
+ evaluated = dict(
+ zip((_hashable(value) for value in distinct), cached, strict=False)
+ )
+ return [evaluated[_hashable(value)] for value in values]
+
+
+def _run_probe(
+ database: Database,
+ catalog: str | None,
+ schema: str | None,
+ transform: str,
+ distinct: list[Any],
+) -> list[Any] | None:
+ try:
+ sql = build_probe_sql(transform, distinct, _dialect_for(database))
+ frame = database.get_df(sql=sql, catalog=catalog, schema=schema)
+ if frame is None or frame.empty:
+ logger.warning(
+ "Partition transform probe returned no rows; skipping
mirroring"
+ )
+ return None
+ row = frame.iloc[0]
+ if len(row) < len(distinct):
+ # The results cannot be aligned back to their inputs; skipping
+ # beats guessing which value produced which column.
+ logger.warning(
+ "Partition transform probe returned %d values for %d inputs",
+ len(row),
+ len(distinct),
+ )
+ return None
+ return [row.iloc[index] for index in range(len(distinct))]
+ except Exception: # pylint: disable=broad-except
+ logger.warning(
+ "Partition transform probe failed; queries will not prune",
+ exc_info=True,
+ )
+ return None
+
+
+def _probe_cache_key(
+ database: Database,
+ catalog: str | None,
+ schema: str | None,
+ transform: str,
+ values: list[Any],
+) -> str:
+ """
+ Key on everything that can change the answer.
+
+ Note this cache is independent of the chart-data cache: it is keyed on the
+ transform and its inputs, so it is correct to share across every chart on
+ every dataset that happens to use the same transform.
+ """
+ payload = json.dumps(
+ [
+ database.id,
+ database.backend,
+ catalog,
+ schema,
+ transform,
+ [repr(value) for value in values],
+ ],
Review Comment:
✅ **CodeAnt verified this suggestion was addressed in subsequent commits and
marked this thread resolved** as of `da9f21c`.
The cache key now includes `database.sqlalchemy_uri`, `database.extra`, and
`database.changed_on` alongside the transform inputs, invalidating entries when
connection or session settings change.
<sub>If that's not right, unresolve this thread and CodeAnt will leave it
open.</sub>
<!-- codeant-auto-resolve-reply -->
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
