diff --git a/scripts/ci/prek/validate_operators_init.py b/scripts/ci/prek/validate_operators_init.py index 0ac6ae5c15c67..b8893700c9037 100755 --- a/scripts/ci/prek/validate_operators_init.py +++ b/scripts/ci/prek/validate_operators_init.py @@ -25,34 +25,82 @@ import ast import sys +from pathlib import Path from typing import Any from rich.console import Console +from rich.markup import escape console = Console(color_system="standard", width=200) -BASE_OPERATOR_CLASS_NAME = "BaseOperator" +# Pre-existing violations exempted from the checks; burn-down tracked at +# https://github.com/apache/airflow/issues/70296 +EXEMPTIONS_PATH = Path(__file__).parent / "validate_operators_init_exemptions.txt" +BASE_CLASS_NAME_SUFFIXES = ("BaseOperator", "BaseSensorOperator") +# Helper callables used as template_fields values, mapped to the fields the helper injects +# on top of the explicit arguments. Injected fields are owned and assigned by the base class, +# so only the explicit arguments are fields the defining class must assign itself. +TEMPLATE_FIELD_HELPERS: dict[str, frozenset[str]] = { + "aws_template_fields": frozenset({"aws_conn_id", "region_name", "verify"}), +} + + +def _resolve_base_name(base: ast.expr) -> str: + """ + Resolve a base-class expression to its plain name. + + Unwraps subscripted generics (``AwsBaseOperator[EmrHook]``) and attribute access + (``module.BaseOperator``). + + :param base: The base-class expression node. + :return: The resolved name, or an empty string when it cannot be resolved. + """ + if isinstance(base, ast.Subscript): + base = base.value + if isinstance(base, ast.Attribute): + return base.attr + if isinstance(base, ast.Name): + return base.id + return "" def _is_operator(class_node: ast.ClassDef) -> bool: """ - Check if a given class node is an operator, based of the string suffix of the base IDs - (ends with "BaseOperator"). + Check if a given class node is an operator or sensor, based on the string suffix of the + base IDs (ends with "BaseOperator" or "BaseSensorOperator"). TODO: Enhance this function to work with nested inheritance trees through dynamic imports. :param class_node: The class node to check. :return: True if the class definition is of an operator, False otherwise. """ - for base in class_node.bases: - if isinstance(base, ast.Name) and base.id.endswith(BASE_OPERATOR_CLASS_NAME): - return True - return False + return any(_resolve_base_name(base).endswith(BASE_CLASS_NAME_SUFFIXES) for base in class_node.bases) + + +def _extract_field_names(value: ast.expr | None) -> list[str] | None: + """ + Extract template-field names from a ``template_fields`` value expression. + + Supports a tuple of constants and known helper calls with constant arguments + (e.g. ``aws_template_fields("s3_bucket", "s3_key")``). For helper calls, fields the + helper injects on behalf of the base class are excluded — see ``TEMPLATE_FIELD_HELPERS``. + + :param value: The value expression assigned to ``template_fields``. + :return: The extracted field names, or None if the expression shape is not supported. + """ + if isinstance(value, ast.Tuple): + return [str(elt.value) for elt in value.elts if isinstance(elt, ast.Constant)] + if isinstance(value, ast.Call): + injected = TEMPLATE_FIELD_HELPERS.get(_resolve_base_name(value.func)) + if injected is not None: + args = [str(arg.value) for arg in value.args if isinstance(arg, ast.Constant)] + return [arg for arg in args if arg not in injected] + return None def _extract_template_fields(class_node: ast.ClassDef) -> list[str]: """ This method takes a class node as input and extracts the template fields from it. Template fields are identified by an assignment statement where the target is a variable - named "template_fields" and the value is a tuple of constants. + named "template_fields" and the value is a tuple of constants or a known helper call. :param class_node: The class node representing the class for which template fields need to be extracted. :return: A list of template fields extracted from the class node. @@ -60,19 +108,15 @@ def _extract_template_fields(class_node: ast.ClassDef) -> list[str]: for class_item in class_node.body: if isinstance(class_item, ast.Assign): for target in class_item.targets: - if ( - isinstance(target, ast.Name) - and target.id == "template_fields" - and isinstance(class_item.value, ast.Tuple) - ): - return [str(elt.value) for elt in class_item.value.elts if isinstance(elt, ast.Constant)] + if isinstance(target, ast.Name) and target.id == "template_fields": + fields = _extract_field_names(class_item.value) + if fields is not None: + return fields elif isinstance(class_item, ast.AnnAssign): - if ( - isinstance(class_item.target, ast.Name) - and class_item.target.id == "template_fields" - and isinstance(class_item.value, ast.Tuple) - ): - return [str(elt.value) for elt in class_item.value.elts if isinstance(elt, ast.Constant)] + if isinstance(class_item.target, ast.Name) and class_item.target.id == "template_fields": + fields = _extract_field_names(class_item.value) + if fields is not None: + return fields return [] @@ -136,7 +180,13 @@ def _handle_constructor_statement( if isinstance(ctor_stmt.targets[0], ast.Attribute): for target in ctor_stmt.targets: if isinstance(target, ast.Attribute) and target.attr in template_fields: - if isinstance(ctor_stmt.value, ast.BoolOp) and isinstance(ctor_stmt.value.op, ast.Or): + if isinstance(ctor_stmt.value, ast.IfExp) and _is_value_preserving_ternary( + ctor_stmt.value, target.attr + ): + _handle_assigned_field( + assigned_template_fields, invalid_assignments, target, ctor_stmt.value.body + ) + elif isinstance(ctor_stmt.value, ast.BoolOp) and isinstance(ctor_stmt.value.op, ast.Or): _handle_assigned_field( assigned_template_fields, invalid_assignments, target, ctor_stmt.value.values[0] ) @@ -173,6 +223,175 @@ def _handle_assigned_field( assigned_template_fields.append(target.attr) +def _target_name(target: ast.expr) -> str | None: + """ + Resolve an assignment target to the field name it binds. + + :param target: The assignment target node. + :return: The attribute name for ``self.`` targets, the identifier for bare-name + targets, or None for anything else. + """ + if isinstance(target, ast.Attribute) and isinstance(target.value, ast.Name) and target.value.id == "self": + return target.attr + if isinstance(target, ast.Name): + return target.id + return None + + +def _is_super_init_call(node: ast.Call) -> bool: + """ + Check whether a call node is ``super().__init__(...)``. + + :param node: The call node to check. + :return: True if the node calls ``__init__`` on a ``super()`` call. + """ + return ( + isinstance(node.func, ast.Attribute) + and node.func.attr == "__init__" + and isinstance(node.func.value, ast.Call) + and isinstance(node.func.value.func, ast.Name) + and node.func.value.func.id == "super" + ) + + +def _is_value_preserving_ternary(value: ast.IfExp, field: str) -> bool: + """ + Check whether a ternary keeps the field value intact when it is set. + + Matches ``field if field else `` and ``field if field is not None else `` + — both equivalent to the sanctioned ``field or `` defaulting idiom. + + :param value: The ternary expression node. + :param field: The template field name. + :return: True if the ternary only substitutes a default for an unset value. + """ + if not (isinstance(value.body, ast.Name) and value.body.id == field): + return False + test = value.test + if isinstance(test, ast.Name) and test.id == field: + return True + return ( + isinstance(test, ast.Compare) + and isinstance(test.left, ast.Name) + and test.left.id == field + and len(test.ops) == 1 + and isinstance(test.ops[0], (ast.Is, ast.IsNot)) + and len(test.comparators) == 1 + and isinstance(test.comparators[0], ast.Constant) + and test.comparators[0].value is None + ) + + +def _collect_sanctioned_uses(ctor: ast.FunctionDef, template_fields: list[str]) -> set[int]: + """ + Collect the AST node ids of template-field reads that belong to sanctioned patterns. + + Sanctioned patterns are the ones the project documents as safe in a constructor: + ``self.field = field``, ``self.field = field or ``, the equivalent + value-preserving ternaries, the local rebind ``field = field or ``, + tuple assignments pairing names one-to-one, and forwarding via + ``super().__init__(field=field)``. + + :param ctor: The constructor function node. + :param template_fields: The template fields of the class. + :return: Set of ``id()``s of Name nodes participating in sanctioned patterns. + """ + sanctioned: set[int] = set() + + def mark(value: ast.expr | None, field: str) -> None: + if isinstance(value, ast.BoolOp) and isinstance(value.op, ast.Or): + value = value.values[0] + elif isinstance(value, ast.IfExp) and _is_value_preserving_ternary(value, field): + for name_node in ast.walk(value.test): + if isinstance(name_node, ast.Name) and name_node.id == field: + sanctioned.add(id(name_node)) + value = value.body + if isinstance(value, ast.Name) and value.id == field: + sanctioned.add(id(value)) + + for node in ast.walk(ctor): + if isinstance(node, (ast.Assign, ast.AnnAssign)): + targets = node.targets if isinstance(node, ast.Assign) else [node.target] + pairs: list[tuple[ast.expr, ast.expr | None]] + if len(targets) == 1 and isinstance(targets[0], ast.Tuple) and isinstance(node.value, ast.Tuple): + pairs = list(zip(targets[0].elts, node.value.elts)) + else: + pairs = [(target, node.value) for target in targets] + for target, value in pairs: + name = _target_name(target) + if name is not None and name in template_fields: + mark(value, name) + elif isinstance(node, ast.Call) and _is_super_init_call(node): + for keyword in node.keywords: + if keyword.arg is not None and keyword.arg in template_fields: + mark(keyword.value, keyword.arg) + return sanctioned + + +def _check_constructor_field_logic( + class_node: ast.ClassDef, template_fields: list[str], source_lines: list[str] +) -> int: + """ + Check a class's constructor for logic applied to template fields. + + Template fields are rendered after the constructor runs, so any read of a template-field + parameter (or ``self.``) outside the sanctioned assignment/forwarding patterns — + validation calls, conditionals, transformations, string interpolation — operates on the + un-rendered Jinja expression and must move to ``execute()``. + + :param class_node: The AST node representing the class definition. + :param template_fields: The template fields of the class. + :param source_lines: The source lines of the file, for reporting. + :return: The number of offending source lines found. + """ + ctor = next( + (item for item in class_node.body if isinstance(item, ast.FunctionDef) and item.name == "__init__"), + None, + ) + if ctor is None or not template_fields: + return 0 + sanctioned = _collect_sanctioned_uses(ctor, template_fields) + args = ctor.args + # Only names bound in the constructor scope can refer to a template field; without this, + # a field named e.g. "json" would false-positive on uses of the stdlib module. + bound_names = {arg.arg for arg in [*args.posonlyargs, *args.args, *args.kwonlyargs]} + bound_names |= { + node.id for node in ast.walk(ctor) if isinstance(node, ast.Name) and isinstance(node.ctx, ast.Store) + } + # Parameter defaults evaluate at class-definition scope, where a name that matches a + # template field (e.g. a field named "conf" vs. the configuration module) is not the field. + in_defaults = { + id(node) + for default in [*args.defaults, *args.kw_defaults] + if default is not None + for node in ast.walk(default) + } + + findings: dict[int, set[str]] = {} + for node in ast.walk(ctor): + if id(node) in in_defaults: + continue + if isinstance(node, ast.Name) and isinstance(node.ctx, ast.Load): + if node.id in template_fields and node.id in bound_names and id(node) not in sanctioned: + findings.setdefault(node.lineno, set()).add(node.id) + elif isinstance(node, ast.Attribute) and isinstance(node.ctx, ast.Load): + if isinstance(node.value, ast.Name) and node.value.id == "self" and node.attr in template_fields: + findings.setdefault(node.lineno, set()).add(f"self.{node.attr}") + + if findings: + console.print( + f"{class_node.name}'s constructor applies logic to template fields. Template fields " + f"are rendered after the constructor runs, so validation or transformation here acts " + f"on the un-rendered Jinja expression and should move to execute():" + ) + for lineno in sorted(findings): + source = source_lines[lineno - 1].strip() if lineno <= len(source_lines) else "" + console.print( + f"[red] line {lineno}: {escape(source)} ({', '.join(sorted(findings[lineno]))})[/red]" + ) + return len(findings) + + def _check_constructor_template_fields(class_node: ast.ClassDef, template_fields: list[str]) -> int: """ This method checks a class's constructor for missing or invalid assignments of template fields. @@ -220,6 +439,47 @@ def _check_constructor_template_fields(class_node: ast.ClassDef, template_fields return count +def _load_exemptions() -> dict[str, set[str]]: + """ + Load the exemption list for known violations that predate the constructor-logic check. + + Each non-comment line has the form ``::``. Exempted classes + are skipped by all checks; an exempted class with no findings fails as stale so the entry + is removed in the same PR that fixes the class, until the list is empty. + + :return: Mapping of repo-relative file path to the exempted class names in that file. + """ + exemptions: dict[str, set[str]] = {} + if not EXEMPTIONS_PATH.exists(): + return exemptions + for raw_line in EXEMPTIONS_PATH.read_text().splitlines(): + line = raw_line.strip() + if not line or line.startswith("#"): + continue + entry_path, sep, class_name = line.partition("::") + if sep and class_name: + exemptions.setdefault(entry_path, set()).add(class_name) + return exemptions + + +def _get_exempted_classes(path: str, exemptions: dict[str, set[str]]) -> set[str]: + """ + Find the exempted class names for a checked file. + + Exemption entries store repo-relative paths while the hook may receive paths relative + to another working directory, so entries are matched as path suffixes. + + :param path: The file path as passed to the script. + :param exemptions: The loaded exemption mapping. + :return: The exempted class names for this file, or an empty set. + """ + resolved = Path(path).resolve().as_posix() + for entry_path, classes in exemptions.items(): + if resolved == entry_path or resolved.endswith(f"/{entry_path}"): + return classes + return set() + + def main(): """ Check missing or invalid template fields in constructors of providers' operators. @@ -227,15 +487,36 @@ def main(): :return: The total number of errors found. """ err = 0 + exemptions = _load_exemptions() for path in sys.argv[1:]: console.print(f"[yellow]{path}[/yellow]") - tree = ast.parse(open(path).read()) + source = open(path).read() + source_lines = source.splitlines() + tree = ast.parse(source) + exempted_classes = _get_exempted_classes(path, exemptions) + exempted_finding_counts: dict[str, int] = {} for node in ast.walk(tree): if isinstance(node, ast.ClassDef) and _is_operator(class_node=node): template_fields = _extract_template_fields(node) or [] + if node.name in exempted_classes: + with console.capture(): + count = _check_constructor_template_fields(node, template_fields) + count += _check_constructor_field_logic(node, template_fields, source_lines) + exempted_finding_counts[node.name] = count + continue err += _check_constructor_template_fields(node, template_fields) + err += _check_constructor_field_logic(node, template_fields, source_lines) + for class_name in sorted(exempted_classes): + if not exempted_finding_counts.get(class_name): + err += 1 + console.print( + f"[red]Stale exemption for {class_name} — the class has no findings anymore " + f"(or is not detected as an operator); remove its entry from " + f"{EXEMPTIONS_PATH.name}[/red]" + ) return err if __name__ == "__main__": - sys.exit(main()) + # A raw error count wraps at 256 (e.g. 256 findings -> exit code 0), so clamp to 0/1. + sys.exit(1 if main() else 0) diff --git a/scripts/ci/prek/validate_operators_init_exemptions.txt b/scripts/ci/prek/validate_operators_init_exemptions.txt new file mode 100644 index 0000000000000..01fd5bc56dbb7 --- /dev/null +++ b/scripts/ci/prek/validate_operators_init_exemptions.txt @@ -0,0 +1,91 @@ +# Known violations that predate the template-field constructor checks, exempted so the +# prek hook can enforce the rule on new code while these are burned down. +# +# Format: :: (one class per line) +# +# Fixing a class (moving template-field validation/transformation out of __init__ into +# execute()) MUST remove its entry in the same PR — the hook fails on stale entries. +# Burn-down tracked at https://github.com/apache/airflow/issues/70296 +providers/amazon/src/airflow/providers/amazon/aws/operators/appflow.py::AppflowBaseOperator +providers/amazon/src/airflow/providers/amazon/aws/operators/bedrock.py::BedrockCreateKnowledgeBaseOperator +providers/amazon/src/airflow/providers/amazon/aws/operators/bedrock.py::BedrockRaGOperator +providers/amazon/src/airflow/providers/amazon/aws/operators/datasync.py::DataSyncOperator +providers/amazon/src/airflow/providers/amazon/aws/operators/dms.py::DmsModifyTaskOperator +providers/amazon/src/airflow/providers/amazon/aws/operators/dms.py::DmsStartReplicationOperator +providers/amazon/src/airflow/providers/amazon/aws/operators/ecs.py::EcsRunTaskOperator +providers/amazon/src/airflow/providers/amazon/aws/operators/emr.py::EmrAddStepsOperator +providers/amazon/src/airflow/providers/amazon/aws/operators/glue.py::GlueDataQualityOperator +providers/amazon/src/airflow/providers/amazon/aws/operators/neptune.py::NeptuneStartDbClusterOperator +providers/amazon/src/airflow/providers/amazon/aws/operators/neptune.py::NeptuneStopDbClusterOperator +providers/amazon/src/airflow/providers/amazon/aws/operators/s3.py::S3DeleteObjectsOperator +providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker.py::SageMakerCreateNotebookOperator +providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker.py::SageMakerProcessingOperator +providers/amazon/src/airflow/providers/amazon/aws/operators/step_function.py::StepFunctionStartExecutionOperator +providers/amazon/src/airflow/providers/amazon/aws/transfers/base.py::AwsToAwsBaseOperator +providers/amazon/src/airflow/providers/amazon/aws/transfers/gcs_to_s3.py::GCSToS3Operator +providers/amazon/src/airflow/providers/amazon/aws/transfers/mongo_to_s3.py::MongoToS3Operator +providers/amazon/src/airflow/providers/amazon/aws/transfers/s3_to_redshift.py::S3ToRedshiftOperator +providers/anthropic/src/airflow/providers/anthropic/operators/agent.py::AnthropicAgentSessionOperator +providers/apache/hive/src/airflow/providers/apache/hive/sensors/hive_partition.py::HivePartitionSensor +providers/apache/hive/src/airflow/providers/apache/hive/sensors/named_hive_partition.py::NamedHivePartitionSensor +providers/apache/kafka/src/airflow/providers/apache/kafka/operators/produce.py::ProduceToTopicOperator +providers/apache/spark/src/airflow/providers/apache/spark/operators/spark_submit.py::SparkSubmitOperator +providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/kueue.py::KubernetesInstallKueueOperator +providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py::KubernetesPodOperator +providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/resource.py::KubernetesResourceBaseOperator +providers/cohere/src/airflow/providers/cohere/operators/embedding.py::CohereEmbeddingOperator +providers/common/ai/src/airflow/providers/common/ai/operators/agent.py::AgentOperator +providers/common/ai/src/airflow/providers/common/ai/operators/document_loader.py::DocumentLoaderOperator +providers/databricks/src/airflow/providers/databricks/operators/databricks_repos.py::DatabricksReposCreateOperator +providers/databricks/src/airflow/providers/databricks/operators/databricks_repos.py::DatabricksReposDeleteOperator +providers/databricks/src/airflow/providers/databricks/operators/databricks_repos.py::DatabricksReposUpdateOperator +providers/databricks/src/airflow/providers/databricks/operators/databricks_sql.py::DatabricksCopyIntoOperator +providers/databricks/src/airflow/providers/databricks/sensors/databricks.py::DatabricksSQLStatementsSensor +providers/dbt/cloud/src/airflow/providers/dbt/cloud/operators/dbt.py::DbtCloudGetJobRunArtifactOperator +providers/docker/src/airflow/providers/docker/operators/docker.py::DockerOperator +providers/google/src/airflow/providers/google/cloud/operators/bigquery.py::BigQueryInsertJobOperator +providers/google/src/airflow/providers/google/cloud/operators/cloud_batch.py::CloudBatchSubmitJobOperator +providers/google/src/airflow/providers/google/cloud/operators/cloud_build.py::CloudBuildCreateBuildOperator +providers/google/src/airflow/providers/google/cloud/operators/cloud_storage_transfer_service.py::CloudDataTransferServiceCreateJobOperator +providers/google/src/airflow/providers/google/cloud/operators/compute.py::ComputeEngineCopyInstanceTemplateOperator +providers/google/src/airflow/providers/google/cloud/operators/compute.py::ComputeEngineDeleteInstanceGroupManagerOperator +providers/google/src/airflow/providers/google/cloud/operators/compute.py::ComputeEngineDeleteInstanceOperator +providers/google/src/airflow/providers/google/cloud/operators/compute.py::ComputeEngineDeleteInstanceTemplateOperator +providers/google/src/airflow/providers/google/cloud/operators/compute.py::ComputeEngineInsertInstanceFromTemplateOperator +providers/google/src/airflow/providers/google/cloud/operators/compute.py::ComputeEngineInsertInstanceGroupManagerOperator +providers/google/src/airflow/providers/google/cloud/operators/compute.py::ComputeEngineInsertInstanceOperator +providers/google/src/airflow/providers/google/cloud/operators/compute.py::ComputeEngineInsertInstanceTemplateOperator +providers/google/src/airflow/providers/google/cloud/operators/compute.py::ComputeEngineInstanceGroupUpdateManagerTemplateOperator +providers/google/src/airflow/providers/google/cloud/operators/compute.py::ComputeEngineSetMachineTypeOperator +providers/google/src/airflow/providers/google/cloud/operators/dataproc.py::DataprocCreateClusterOperator +providers/google/src/airflow/providers/google/cloud/operators/dataproc.py::DataprocSubmitJobOperator +providers/google/src/airflow/providers/google/cloud/operators/functions.py::CloudFunctionDeployFunctionOperator +providers/google/src/airflow/providers/google/cloud/operators/gcs.py::GCSDeleteObjectsOperator +providers/google/src/airflow/providers/google/cloud/operators/gcs.py::GCSFileTransformOperator +providers/google/src/airflow/providers/google/cloud/operators/gcs.py::GCSListObjectsOperator +providers/google/src/airflow/providers/google/cloud/operators/gen_ai.py::GenAIGeminiCreateBatchJobOperator +providers/google/src/airflow/providers/google/cloud/operators/gen_ai.py::GenAIGeminiCreateEmbeddingsBatchJobOperator +providers/google/src/airflow/providers/google/cloud/sensors/bigquery_dts.py::BigQueryDataTransferServiceTransferRunSensor +providers/google/src/airflow/providers/google/cloud/sensors/cloud_composer.py::CloudComposerExternalTaskSensor +providers/google/src/airflow/providers/google/cloud/transfers/azure_fileshare_to_gcs.py::AzureFileShareToGCSOperator +providers/google/src/airflow/providers/google/cloud/transfers/bigquery_to_mssql.py::BigQueryToMsSqlOperator +providers/google/src/airflow/providers/google/cloud/transfers/gcs_to_bigquery.py::GCSToBigQueryOperator +providers/google/src/airflow/providers/google/cloud/transfers/gcs_to_gcs.py::GCSToGCSOperator +providers/google/src/airflow/providers/google/cloud/transfers/gcs_to_local.py::GCSToLocalFilesystemOperator +providers/google/src/airflow/providers/google/marketing_platform/operators/campaign_manager.py::GoogleCampaignManagerDeleteReportOperator +providers/microsoft/azure/src/airflow/providers/microsoft/azure/sensors/compute.py::AzureVirtualMachineStateSensor +providers/microsoft/azure/src/airflow/providers/microsoft/azure/transfers/gcs_to_wasb.py::GCSToAzureBlobStorageOperator +providers/microsoft/azure/src/airflow/providers/microsoft/azure/transfers/oracle_to_azure_data_lake.py::OracleToAzureDataLakeOperator +providers/microsoft/psrp/src/airflow/providers/microsoft/psrp/operators/psrp.py::PsrpOperator +providers/neo4j/src/airflow/providers/neo4j/operators/neo4j.py::Neo4jOperator +providers/oracle/src/airflow/providers/oracle/transfers/oracle_to_oracle.py::OracleToOracleOperator +providers/papermill/src/airflow/providers/papermill/operators/papermill.py::PapermillOperator +providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py::SnowparkContainerJobOperator +providers/ssh/src/airflow/providers/ssh/operators/ssh.py::SSHOperator +providers/ssh/src/airflow/providers/ssh/operators/ssh_remote_job.py::SSHRemoteJobOperator +providers/standard/src/airflow/providers/standard/operators/bash.py::BashOperator +providers/standard/src/airflow/providers/standard/operators/hitl.py::HITLOperator +providers/standard/src/airflow/providers/standard/operators/trigger_dagrun.py::TriggerDagRunOperator +providers/standard/src/airflow/providers/standard/sensors/date_time.py::DateTimeSensor +providers/teradata/src/airflow/providers/teradata/transfers/teradata_to_teradata.py::TeradataToTeradataOperator +providers/weaviate/src/airflow/providers/weaviate/operators/weaviate.py::WeaviateIngestOperator diff --git a/scripts/tests/ci/prek/test_validate_operators_init.py b/scripts/tests/ci/prek/test_validate_operators_init.py new file mode 100644 index 0000000000000..cc308796e4a5c --- /dev/null +++ b/scripts/tests/ci/prek/test_validate_operators_init.py @@ -0,0 +1,206 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import ast +import sys +import textwrap +from pathlib import Path + +import pytest +import validate_operators_init +from validate_operators_init import ( + _check_constructor_field_logic, + _check_constructor_template_fields, + _extract_template_fields, + _is_operator, + main, +) + + +def _first_class(code: str) -> ast.ClassDef: + tree = ast.parse(textwrap.dedent(code)) + return next(node for node in ast.walk(tree) if isinstance(node, ast.ClassDef)) + + +def _logic_findings(code: str, template_fields: list[str]) -> int: + code = textwrap.dedent(code) + return _check_constructor_field_logic(_first_class(code), template_fields, code.splitlines()) + + +def _operator_code(ctor_body: str) -> str: + body = textwrap.indent(textwrap.dedent(ctor_body).strip("\n"), " " * 8) + return ( + "class MyOperator(BaseOperator):\n" + ' template_fields = ("foo",)\n' + "\n" + " def __init__(self, foo=None, **kwargs):\n" + f"{body}\n" + ) + + +class TestConstructorFieldLogic: + @pytest.mark.parametrize( + "ctor_body, expected", + [ + pytest.param("self.foo = foo", 0, id="plain-assignment"), + pytest.param("self.foo = foo or 'default'", 0, id="or-default"), + pytest.param("self.foo = foo if foo else 'default'", 0, id="ternary-truthiness"), + pytest.param("self.foo = foo if foo is not None else 'default'", 0, id="ternary-is-not-none"), + pytest.param("foo = foo or 'default'\nself.foo = foo", 0, id="local-rebind"), + pytest.param("super().__init__(foo=foo)", 0, id="super-forwarding"), + pytest.param("self._validate(foo)\nself.foo = foo", 1, id="validation-call"), + pytest.param( + "if foo is not None:\n self._validate(foo)\nself.foo = foo", + 2, + id="nested-validation-call", + ), + pytest.param("self.foo = foo\nself.bar = foo.upper()", 1, id="derived-assignment"), + pytest.param( + "self.foo = foo\nif self.foo:\n self.bar = 1", + 1, + id="self-attribute-read", + ), + pytest.param( + "if not foo:\n raise ValueError(f'unsupported: {foo}')\nself.foo = foo", + 2, + id="ctor-validation-raise", + ), + ], + ) + def test_flags_logic_but_not_sanctioned_patterns(self, ctor_body: str, expected: int): + assert _logic_findings(_operator_code(ctor_body), ["foo"]) == expected + + def test_name_in_parameter_default_is_not_the_field(self): + # A field named like a module (e.g. "conf") used in a parameter default evaluates at + # class-definition scope and must not be flagged. + code = """ + class MyOperator(BaseOperator): + template_fields = ("conf",) + + def __init__(self, conf=None, deferrable=conf.getboolean("operators", "x"), **kwargs): + self.conf = conf + """ + assert _logic_findings(code, ["conf"]) == 0 + + def test_unbound_module_name_matching_field_is_not_flagged(self): + code = """ + class MyOperator(BaseOperator): + template_fields = ("json",) + + def __init__(self, **kwargs): + self.data = json.dumps({}) + """ + assert _logic_findings(code, ["json"]) == 0 + + def test_no_constructor_is_clean(self): + code = """ + class MyOperator(BaseOperator): + template_fields = ("foo",) + """ + assert _logic_findings(code, ["foo"]) == 0 + + +class TestOperatorDetection: + @pytest.mark.parametrize( + "class_def, expected", + [ + pytest.param("class Op(BaseOperator):", True, id="base-operator"), + pytest.param("class Sensor(BaseSensorOperator):", True, id="base-sensor"), + pytest.param("class Op(AwsBaseOperator[EmrHook]):", True, id="subscripted-base"), + pytest.param("class Helper:", False, id="plain-class"), + pytest.param("class Hook(BaseHook):", False, id="hook"), + ], + ) + def test_detects_operator_bases(self, class_def: str, expected: bool): + assert _is_operator(_first_class(f"{class_def}\n pass")) is expected + + +class TestTemplateFieldExtraction: + def test_extracts_helper_call_without_injected_fields(self): + code = """ + class Op(AwsBaseOperator[EC2Hook]): + template_fields = aws_template_fields("instance_id", "region_name", "aws_conn_id") + """ + # region_name / aws_conn_id are injected and assigned by the AWS base operator. + assert _extract_template_fields(_first_class(code)) == ["instance_id"] + + def test_extracts_tuple(self): + code = """ + class Op(BaseOperator): + template_fields = ("a", "b") + """ + assert _extract_template_fields(_first_class(code)) == ["a", "b"] + + +class TestValuePreservingTernaryAssignment: + def test_ternary_default_is_a_valid_assignment(self): + code = """ + class Op(BaseOperator): + template_fields = ("conf",) + + def __init__(self, conf=None, **kwargs): + self.conf = conf if conf else {} + """ + assert _check_constructor_template_fields(_first_class(code), ["conf"]) == 0 + + +class TestExemptions: + VIOLATING_OPERATOR = textwrap.dedent( + """ + class MyOperator(BaseOperator): + template_fields = ("foo",) + + def __init__(self, foo=None, **kwargs): + self._validate(foo) + self.foo = foo + """ + ) + CLEAN_OPERATOR = textwrap.dedent( + """ + class MyOperator(BaseOperator): + template_fields = ("foo",) + + def __init__(self, foo=None, **kwargs): + self.foo = foo + """ + ) + + def _run(self, monkeypatch, tmp_path: Path, code: str, exemption_line: str | None) -> int: + target = tmp_path / "my_operator.py" + target.write_text(code) + exemptions = tmp_path / "exemptions.txt" + exemptions.write_text(f"# comment\n{exemption_line}\n" if exemption_line else "# comment\n") + monkeypatch.setattr(validate_operators_init, "EXEMPTIONS_PATH", exemptions) + monkeypatch.setattr(sys, "argv", ["validate_operators_init.py", str(target)]) + return main() + + def test_exempted_class_is_suppressed(self, monkeypatch, tmp_path: Path): + err = self._run(monkeypatch, tmp_path, self.VIOLATING_OPERATOR, "my_operator.py::MyOperator") + assert err == 0 + + def test_violation_without_exemption_fails(self, monkeypatch, tmp_path: Path): + err = self._run(monkeypatch, tmp_path, self.VIOLATING_OPERATOR, None) + assert err > 0 + + def test_stale_exemption_fails(self, monkeypatch, tmp_path: Path): + err = self._run(monkeypatch, tmp_path, self.CLEAN_OPERATOR, "my_operator.py::MyOperator") + assert err == 1 + + def test_exemption_for_other_class_does_not_apply(self, monkeypatch, tmp_path: Path): + err = self._run(monkeypatch, tmp_path, self.VIOLATING_OPERATOR, "my_operator.py::OtherOperator") + assert err > 0