#
# 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 os
import urllib.parse
import warnings
from collections.abc import Sequence
from typing import TYPE_CHECKING, Any, cast
from botocore.exceptions import ClientError
from airflow.exceptions import AirflowProviderDeprecationWarning
from airflow.providers.amazon.aws.exceptions import GlueJobRunStoppedError
from airflow.providers.amazon.aws.hooks.glue import GlueDataQualityHook, GlueJobHook
from airflow.providers.amazon.aws.hooks.s3 import S3Hook
from airflow.providers.amazon.aws.links.glue import GlueJobRunDetailsLink
from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
from airflow.providers.amazon.aws.triggers.glue import (
GlueDataQualityRuleRecommendationRunCompleteTrigger,
GlueDataQualityRuleSetEvaluationRunCompleteTrigger,
GlueJobCompleteTrigger,
)
from airflow.providers.amazon.aws.utils import validate_execute_complete_event
from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
from airflow.providers.amazon.version_compat import AIRFLOW_V_3_3_PLUS
from airflow.providers.common.compat.openlineage.utils.spark import (
inject_parent_job_information_into_glue_arguments,
inject_transport_information_into_glue_arguments,
)
from airflow.providers.common.compat.sdk import AirflowException, conf
_DURABLE_UNSET = object()
def _warn_and_disable_durable_pre_3_3(durable: Any) -> bool:
"""Shared by the <3.3 compat stub: durable has no effect below 3.3, warn if it was set."""
if durable is not _DURABLE_UNSET:
warnings.warn(
"`durable` has no effect on Airflow versions below 3.3.",
UserWarning,
stacklevel=3,
)
return False
# ResumableJobMixin only exists on Airflow 3.3+; this provider still targets >=2.11. Drop this
# fallback once the provider's minimum Airflow version is >=3.3.
try:
from airflow.sdk import ResumableJobMixin
except ImportError:
[docs]
class ResumableJobMixin: # type: ignore[no-redef]
"""Airflow <3.3 stub, task_state_store unavailable, always submits fresh."""
[docs]
external_id_key: str = "glue_job_run_id"
def __init__(self, *, durable: Any = _DURABLE_UNSET, **kwargs: Any) -> None:
super().__init__(**kwargs)
[docs]
self.durable = _warn_and_disable_durable_pre_3_3(durable)
[docs]
def execute_resumable(self, context):
external_id = self.submit_job(context)
self.poll_until_complete(external_id, context)
return self.get_job_result(external_id, context)
if TYPE_CHECKING:
from pydantic import JsonValue
from airflow.sdk import Context
# Glue job run states, see
# https://docs.aws.amazon.com/glue/latest/dg/aws-glue-api-jobs-runs.html#aws-glue-api-jobs-runs-JobRun
# STOPPED is deliberately NOT a success state here, unlike GlueJobHook.job_completion: Glue can't
# tell a console cancellation apart from on_kill stopping its own run, so treating STOPPED as
# failure avoids silently reporting a self-inflicted stop as success.
[docs]
JOB_RUN_SUCCESS_STATES = ("SUCCEEDED",)
[docs]
JOB_RUN_TERMINAL_STATES = (*JOB_RUN_SUCCESS_STATES, "STOPPED", "FAILED", "TIMEOUT", "ERROR", "EXPIRED")
# Synthetic state for a run id Glue no longer knows about, so a retry submits fresh rather than failing.
[docs]
NOT_FOUND_STATE = "NOT_FOUND"
[docs]
class GlueJobOperator(ResumableJobMixin, AwsBaseOperator[GlueJobHook]):
"""
Create an AWS Glue Job.
AWS Glue is a serverless Spark ETL service for running Spark Jobs on the AWS
cloud. Language support: Python and Scala.
.. seealso::
For more information on how to use this operator, take a look at the guide:
:ref:`howto/operator:GlueJobOperator`
:param job_name: unique job name per AWS Account
:param script_location: location of ETL script. Must be a local or S3 path
:param job_desc: job description details
:param concurrent_run_limit: The maximum number of concurrent runs allowed for a job
:param script_args: etl script arguments and AWS Glue arguments (templated)
:param retry_limit: The maximum number of times to retry this job if it fails
:param num_of_dpus: Number of AWS Glue DPUs to allocate to this Job.
:param s3_bucket: S3 bucket where logs and local etl script will be uploaded
:param iam_role_name: AWS IAM Role for Glue Job Execution. If set `iam_role_arn` must equal None.
:param iam_role_arn: AWS IAM ARN for Glue Job Execution. If set `iam_role_name` must equal None.
:param create_job_kwargs: Extra arguments for Glue Job Creation
:param run_job_kwargs: Extra arguments for Glue Job Run
:param wait_for_completion: Whether to wait for job run completion. (default: True)
:param deferrable: If True, the operator will wait asynchronously for the job to complete.
This implies waiting for completion. This mode requires aiobotocore module to be installed.
(default: False)
:param verbose: If True, Glue Job Run logs show in the Airflow Task Logs. (default: False)
:param update_config: If True, Operator will update job configuration. (default: False)
:param replace_script_file: If True, the script file will be replaced in S3. (default: False)
:param stop_job_run_on_kill: If True, Operator will stop the job run when task is killed.
:param sleep_before_return: time in seconds to wait before returning final status. This is meaningful in case
of limiting concurrency, Glue needs 5-10 seconds to clean up resources.
Thus if status is returned immediately it might end up in case of more than 1 concurrent run.
It is recommended to set this parameter to 10 when you are using concurrency=1.
For more information see:
https://repost.aws/questions/QUaKgpLBMPSGWO0iq2Fob_bw/glue-run-concurrent-jobs#ANFpCL2fRnQRqgDFuIU_rpvA
:param openlineage_inject_parent_job_info: If True, injects OpenLineage parent job information into the
Glue job's ``--conf`` argument so the Glue Spark job emits a ``parentRunFacet`` linking back to the
Airflow task. Defaults to the ``openlineage.spark_inject_parent_job_info`` config value.
:param openlineage_inject_transport_info: If True, injects OpenLineage transport configuration into the
Glue job's ``--conf`` argument so the Glue Spark job sends OL events to the same backend as Airflow.
Defaults to the ``openlineage.spark_inject_transport_info`` config value.
:param waiter_delay: Time in seconds to wait between status checks. (default: 60)
:param waiter_max_attempts: Maximum number of attempts to check for job completion. (default: 20)
:param resume_glue_job_on_retry: deprecated, use ``durable`` instead.
:param durable: When ``True``, the Glue job run id is persisted to task state before polling
begins. A worker crash on retry reconnects to the existing run instead of starting a
duplicate. Defaults to ``True`` on Airflow 3.3+, which uses task state store for the
persisted lookup; on earlier versions it defaults to ``False`` and, if set explicitly,
recovers the run by searching job runs for the task's ``--airflow_task_uuid`` argument.
:param aws_conn_id: The Airflow connection used for AWS credentials.
If this is ``None`` or empty then the default boto3 behaviour is used. If
running Airflow in a distributed manner and aws_conn_id is None or
empty, then default boto3 configuration would be used (and must be
maintained on each worker node).
:param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
:param verify: Whether or not to verify SSL certificates. See:
https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
:param botocore_config: Configuration dictionary (key-values) for botocore client. See:
https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
"""
[docs]
aws_hook_class = GlueJobHook
[docs]
template_fields: Sequence[str] = aws_template_fields(
"job_name",
"script_location",
"script_args",
"create_job_kwargs",
"run_job_kwargs",
"s3_bucket",
"iam_role_name",
"iam_role_arn",
)
[docs]
template_ext: Sequence[str] = ()
[docs]
template_fields_renderers = {
"script_args": "json",
"create_job_kwargs": "json",
"run_job_kwargs": "json",
}
[docs]
TASK_UUID_ARG = "--airflow_task_uuid"
[docs]
external_id_key = "glue_job_run_id"
def __init__(
self,
*,
job_name: str = "aws_glue_default_job",
job_desc: str = "AWS Glue Job with Airflow",
script_location: str | None = None,
concurrent_run_limit: int | None = None,
script_args: dict | None = None,
retry_limit: int = 0,
num_of_dpus: int | float | None = None,
s3_bucket: str | None = None,
iam_role_name: str | None = None,
iam_role_arn: str | None = None,
create_job_kwargs: dict | None = None,
run_job_kwargs: dict | None = None,
wait_for_completion: bool = True,
deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
verbose: bool = False,
replace_script_file: bool = False,
update_config: bool = False,
stop_job_run_on_kill: bool = False,
sleep_before_return: int = 0,
job_poll_interval: int | float = 6,
waiter_delay: int = 60,
waiter_max_attempts: int = 75,
resume_glue_job_on_retry: bool | None = None,
durable: bool | None = None,
openlineage_inject_parent_job_info: bool = conf.getboolean(
"openlineage", "spark_inject_parent_job_info", fallback=False
),
openlineage_inject_transport_info: bool = conf.getboolean(
"openlineage", "spark_inject_transport_info", fallback=False
),
**kwargs,
):
if resume_glue_job_on_retry is not None:
warnings.warn(
"`resume_glue_job_on_retry` is deprecated and will be removed once this provider's "
"minimum supported Airflow version reaches 3.3. "
+ (
"Use `durable` instead."
if AIRFLOW_V_3_3_PLUS
else "On Airflow 3.3+, use `durable` instead."
),
AirflowProviderDeprecationWarning,
stacklevel=2,
)
if AIRFLOW_V_3_3_PLUS:
kwargs.setdefault("durable", resume_glue_job_on_retry)
# durable is also named parameter here (not left to **kwargs) so default_args={"durable": ...} reaches
# it on every supported Airflow version.
if durable is not None:
kwargs["durable"] = durable
super().__init__(**kwargs)
if not AIRFLOW_V_3_3_PLUS and resume_glue_job_on_retry is not None:
# durable itself has no effect below 3.3, so we take value of resume_glue_job_on_retry instead.
self.durable = resume_glue_job_on_retry
[docs]
self.job_name = job_name
[docs]
self.job_desc = job_desc
[docs]
self.script_location = script_location
[docs]
self.concurrent_run_limit = concurrent_run_limit or 1
[docs]
self.script_args = script_args or {}
[docs]
self.retry_limit = retry_limit
[docs]
self.num_of_dpus = num_of_dpus
[docs]
self.s3_bucket = s3_bucket
[docs]
self.iam_role_name = iam_role_name
[docs]
self.iam_role_arn = iam_role_arn
[docs]
self.s3_protocol = "s3://"
[docs]
self.s3_artifacts_prefix = "artifacts/glue-scripts/"
[docs]
self.create_job_kwargs = create_job_kwargs
[docs]
self.run_job_kwargs = run_job_kwargs or {}
[docs]
self.wait_for_completion = wait_for_completion
[docs]
self.update_config = update_config
[docs]
self.replace_script_file = replace_script_file
[docs]
self.deferrable = deferrable
[docs]
self.job_poll_interval = job_poll_interval
[docs]
self.stop_job_run_on_kill = stop_job_run_on_kill
self._job_run_id: str | None = None
[docs]
self.sleep_before_return: int = sleep_before_return
[docs]
self.s3_script_location: str | None = None
[docs]
self.waiter_delay = waiter_delay
[docs]
self.waiter_max_attempts = waiter_max_attempts
[docs]
self.openlineage_inject_parent_job_info = openlineage_inject_parent_job_info
[docs]
self.openlineage_inject_transport_info = openlineage_inject_transport_info
@property
def _hook_parameters(self):
# Upload script to S3 before creating the hook.
if self.script_location is None:
self.s3_script_location = None
# location provided, but it's not in S3 yet.
elif self.script_location and self.s3_script_location is None:
if not self.script_location.startswith(self.s3_protocol):
self.upload_etl_script_to_s3()
else:
self.s3_script_location = self.script_location
return {
**super()._hook_parameters,
"job_name": self.job_name,
"desc": self.job_desc,
"concurrent_run_limit": self.concurrent_run_limit,
"script_location": self.s3_script_location,
"retry_limit": self.retry_limit,
"num_of_dpus": self.num_of_dpus,
"aws_conn_id": self.aws_conn_id,
"region_name": self.region_name,
"s3_bucket": self.s3_bucket,
"iam_role_name": self.iam_role_name,
"iam_role_arn": self.iam_role_arn,
"create_job_kwargs": self.create_job_kwargs,
"update_config": self.update_config,
"job_poll_interval": self.job_poll_interval,
}
[docs]
def upload_etl_script_to_s3(self):
"""Upload the ETL script to S3."""
s3_hook = S3Hook(aws_conn_id=self.aws_conn_id)
script_name = os.path.basename(self.script_location)
s3_hook.load_file(
self.script_location,
self.s3_artifacts_prefix + script_name,
bucket_name=self.s3_bucket,
replace=self.replace_script_file,
)
self.s3_script_location = f"s3://{self.s3_bucket}/{self.s3_artifacts_prefix}{script_name}"
def _get_task_uuid(self, context: Context) -> str:
ti = context["ti"]
map_index = getattr(ti, "map_index", -1)
if map_index is None:
map_index = -1
return f"{ti.dag_id}:{ti.task_id}:{ti.run_id}:{map_index}"
def _prepare_script_args_with_task_uuid(
self, context: Context, base_args: dict | None = None
) -> tuple[dict, str]:
script_args = dict(base_args if base_args is not None else (self.script_args or {}))
if self.TASK_UUID_ARG in script_args:
task_uuid = str(script_args[self.TASK_UUID_ARG])
else:
task_uuid = self._get_task_uuid(context)
script_args[self.TASK_UUID_ARG] = task_uuid
return script_args, task_uuid
def _find_job_run_id_by_task_uuid(self, task_uuid: str) -> tuple[str, str] | None:
# Unbounded walk with no page cap; a no-match run is the common retry shape and pays the
# full scan. Tracked at https://github.com/apache/airflow/issues/71489.
next_token: str | None = None
while True:
request = {"JobName": self.job_name, "MaxResults": 50}
if next_token:
request["NextToken"] = next_token
response = self.hook.conn.get_job_runs(**request)
for job_run in response.get("JobRuns", []):
args = job_run.get("Arguments", {}) or {}
if args.get(self.TASK_UUID_ARG) == task_uuid:
job_run_id = job_run.get("Id")
job_run_state = job_run.get("JobRunState")
if job_run_id and job_run_state:
return job_run_id, job_run_state
next_token = response.get("NextToken")
if not next_token:
return None
[docs]
def execute(self, context: Context) -> str | None:
"""
Execute AWS Glue Job from Airflow.
:return: the current Glue job ID.
"""
if self.deferrable:
# The Triggerer tracks the run, so no id is persisted to task_state_store here -- but
# durable still reattaches a retry via the task-UUID scan in submit_job.
job_run_id = self.submit_job(context)
self.defer(
trigger=GlueJobCompleteTrigger(
job_name=self.job_name,
run_id=job_run_id,
verbose=self.verbose,
aws_conn_id=self.aws_conn_id,
waiter_delay=self.waiter_delay,
waiter_max_attempts=self.waiter_max_attempts,
region_name=self.region_name,
),
method_name="execute_complete",
)
self.execute_resumable(context)
return self._job_run_id
[docs]
def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> str:
validated_event = validate_execute_complete_event(event)
if validated_event["status"] != "success":
raise AirflowException(f"Error in glue job: {validated_event}")
return validated_event["run_id"]
[docs]
def on_kill(self):
"""Cancel the running AWS Glue Job."""
if self.stop_job_run_on_kill and self._job_run_id:
self.log.info("Stopping AWS Glue Job: %s. Run Id: %s", self.job_name, self._job_run_id)
response = self.hook.conn.batch_stop_job_run(
JobName=self.job_name,
JobRunIds=[self._job_run_id],
)
if not response["SuccessfulSubmissions"]:
self.log.error("Failed to stop AWS Glue Job: %s. Run Id: %s", self.job_name, self._job_run_id)
def _set_job_run_id(self, context: Context, job_run_id: str) -> None:
"""Record the run id and surface its console link; guarded so one attempt logs it once."""
if self._job_run_id == job_run_id:
return
self._job_run_id = job_run_id
context["ti"].xcom_push(key="glue_job_run_id", value=job_run_id)
GlueJobRunDetailsLink.persist(
context=context,
operator=self,
region_name=self.hook.conn_region_name,
aws_partition=self.hook.conn_partition,
job_name=urllib.parse.quote(self.job_name, safe=""),
job_run_id=job_run_id,
)
self.log.info(
"You can monitor this Glue Job run at: %s",
GlueJobRunDetailsLink.format_str.format(
aws_domain=GlueJobRunDetailsLink.get_aws_domain(self.hook.conn_partition),
region_name=self.hook.conn_region_name,
job_name=urllib.parse.quote(self.job_name, safe=""),
job_run_id=job_run_id,
),
)
def _build_script_args(self, context: Context) -> dict:
script_args = dict(self.script_args)
if self.openlineage_inject_parent_job_info:
self.log.info("Injecting OpenLineage parent job information into Glue job arguments.")
script_args = inject_parent_job_information_into_glue_arguments(script_args, context)
if self.openlineage_inject_transport_info:
self.log.info("Injecting OpenLineage transport information into Glue job arguments.")
script_args = inject_transport_information_into_glue_arguments(script_args, context)
# Skipped for a synchronous run on 3.3+: task_state_store is the primary reconnect
# mechanism there, so tagging every script's args is not worth it just for the narrow
# crash-before-persist window this tag would otherwise cover.
if self.durable and (self.deferrable or not AIRFLOW_V_3_3_PLUS):
script_args, _ = self._prepare_script_args_with_task_uuid(context, base_args=script_args)
return script_args
def _find_previous_job_run(self, context: Context, task_uuid: str) -> str | None:
"""
Look for a Glue job run this task instance already started.
Checks XCom for a cached run id first, then falls back to a task-UUID scan. The XCom tier
only works on Airflow 2.x; every Airflow 3 release clears task XComs before each
non-deferral attempt, so it always misses there and the scan runs every time.
"""
ti = context["ti"]
previous_job_run_id = ti.xcom_pull(key="glue_job_run_id", task_ids=ti.task_id)
if previous_job_run_id:
try:
job_run = self.hook.conn.get_job_run(JobName=self.job_name, RunId=previous_job_run_id)
state = job_run.get("JobRun", {}).get("JobRunState")
self.log.info("Previous Glue job_run_id: %s, state: %s", previous_job_run_id, state)
if self.is_job_active(state):
return previous_job_run_id
except Exception:
self.log.warning("Failed to get previous Glue job run state", exc_info=True)
else:
try:
existing = self._find_job_run_id_by_task_uuid(task_uuid)
if existing:
existing_job_run_id, existing_job_run_state = existing
self.log.info(
"Found Glue job_run_id by task UUID: %s, state: %s",
existing_job_run_id,
existing_job_run_state,
)
if self.is_job_active(existing_job_run_state):
return existing_job_run_id
except Exception:
self.log.warning("Failed to find previous Glue job run by task UUID", exc_info=True)
return None
def _has_stored_external_id(self, context: Context) -> bool:
task_state_store = context.get("task_state_store")
return task_state_store is not None and task_state_store.get(self.external_id_key) is not None
[docs]
def submit_job(self, context: Context) -> str:
"""Start a Glue job run and return its run id, or reconnect to one this task already started."""
script_args = self._build_script_args(context)
# Scan only when there's nothing else to go on: first attempt, no store, or a store that
# never recorded this key. The store check itself is skipped on the deferrable path, where
# nothing is ever written to it -- a store error there would otherwise block the scan.
# Deferrable retries could reconnect via task_state_store too, avoiding this scan entirely,
# if ResumableJobMixin exposed its reconnect decision apart from its polling loop;
# tracked at https://github.com/apache/airflow/issues/71485.
if (
self.durable
and (self.deferrable or not AIRFLOW_V_3_3_PLUS)
and context["ti"].try_number > 1
and (self.deferrable or not self._has_stored_external_id(context))
):
existing_job_run_id = self._find_previous_job_run(context, script_args[self.TASK_UUID_ARG])
if existing_job_run_id:
self._set_job_run_id(context, existing_job_run_id)
return existing_job_run_id
self.log.info(
"Initializing AWS Glue Job: %s. Wait for completion: %s",
self.job_name,
self.wait_for_completion,
)
# A prior get_job_status call may have set this to a stale, terminal run id while checking
# whether to reconnect. Clear it so on_kill has nothing to act on if initialize_job raises.
self._job_run_id = None
glue_job_run = self.hook.initialize_job(script_args, self.run_job_kwargs)
# Set before polling so on_kill can stop the run even if the worker dies immediately after.
self._set_job_run_id(context, glue_job_run["JobRunId"])
return glue_job_run["JobRunId"]
[docs]
def get_job_status(self, external_id: JsonValue, context: Context) -> str:
"""Query the raw job run state; a run id Glue no longer knows about degrades to NOT_FOUND."""
job_run_id = cast("str", external_id)
# This is the first place a reconnecting attempt learns the run id.
self._set_job_run_id(context, job_run_id)
try:
return self.hook.get_job_state(self.job_name, job_run_id)
except ClientError as e:
if e.response["Error"]["Code"] == "EntityNotFoundException":
return NOT_FOUND_STATE
raise
[docs]
def is_job_active(self, status: str) -> bool:
return status not in (*JOB_RUN_TERMINAL_STATES, NOT_FOUND_STATE)
[docs]
def is_job_succeeded(self, status: str) -> bool:
return status in JOB_RUN_SUCCESS_STATES
[docs]
def poll_until_complete(self, external_id: JsonValue, context: Context) -> None:
job_run_id = cast("str", external_id)
self._set_job_run_id(context, job_run_id)
if not self.wait_for_completion:
self.log.info("AWS Glue Job: %s. Run Id: %s", self.job_name, job_run_id)
return
glue_job_run = self.hook.job_completion(
self.job_name, job_run_id, self.verbose, self.sleep_before_return
)
state = glue_job_run["JobRunState"]
self.log.info("AWS Glue Job: %s status: %s. Run Id: %s", self.job_name, state, job_run_id)
if state not in JOB_RUN_SUCCESS_STATES:
# job_completion's own finished_states also accepts STOPPED; this is narrower on purpose.
raise GlueJobRunStoppedError(
f"Glue job run {job_run_id} for job {self.job_name} ended in state {state!r} "
"instead of succeeding."
)
[docs]
def get_job_result(self, external_id: JsonValue, context: Context) -> str:
job_run_id = cast("str", external_id)
self._set_job_run_id(context, job_run_id)
return job_run_id
[docs]
class GlueDataQualityOperator(AwsBaseOperator[GlueDataQualityHook]):
"""
Creates a data quality ruleset with DQDL rules applied to a specified Glue table.
.. seealso::
For more information on how to use this operator, take a look at the guide:
:ref:`howto/operator:GlueDataQualityOperator`
:param name: A unique name for the data quality ruleset.
:param ruleset: A Data Quality Definition Language (DQDL) ruleset.
For more information, see the Glue developer guide.
:param description: A description of the data quality ruleset.
:param update_rule_set: To update existing ruleset, Set this flag to True. (default: False)
:param data_quality_ruleset_kwargs: Extra arguments for RuleSet.
:param aws_conn_id: The Airflow connection used for AWS credentials.
If this is ``None`` or empty then the default boto3 behaviour is used. If
running Airflow in a distributed manner and aws_conn_id is None or
empty, then default boto3 configuration would be used (and must be
maintained on each worker node).
:param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
:param verify: Whether or not to verify SSL certificates. See:
https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
:param botocore_config: Configuration dictionary (key-values) for botocore client. See:
https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
"""
[docs]
aws_hook_class = GlueDataQualityHook
[docs]
template_fields: Sequence[str] = aws_template_fields(
"name", "ruleset", "description", "data_quality_ruleset_kwargs"
)
[docs]
template_fields_renderers = {
"data_quality_ruleset_kwargs": "json",
}
def __init__(
self,
*,
name: str,
ruleset: str,
description: str = "AWS Glue Data Quality Rule Set With Airflow",
update_rule_set: bool = False,
data_quality_ruleset_kwargs: dict | None = None,
**kwargs,
):
super().__init__(**kwargs)
[docs]
self.description = description
[docs]
self.update_rule_set = update_rule_set
[docs]
self.data_quality_ruleset_kwargs = data_quality_ruleset_kwargs or {}
[docs]
def execute(self, context: Context):
# ruleset is a template field; strip the rendered value here, not in __init__.
self.ruleset = self.ruleset.strip()
self.validate_inputs()
config = {
"Name": self.name,
"Ruleset": self.ruleset,
"Description": self.description,
**self.data_quality_ruleset_kwargs,
}
try:
if self.update_rule_set:
self.hook.conn.update_data_quality_ruleset(**config)
self.log.info("AWS Glue data quality ruleset updated successfully")
else:
self.hook.conn.create_data_quality_ruleset(**config)
self.log.info("AWS Glue data quality ruleset created successfully")
except ClientError as error:
raise AirflowException(
f"AWS Glue data quality ruleset failed: {error.response['Error']['Message']}"
)
[docs]
class GlueDataQualityRuleSetEvaluationRunOperator(AwsBaseOperator[GlueDataQualityHook]):
"""
Evaluate a ruleset against a data source (Glue table).
.. seealso::
For more information on how to use this operator, take a look at the guide:
:ref:`howto/operator:GlueDataQualityRuleSetEvaluationRunOperator`
:param datasource: The data source (Glue table) associated with this run. (templated)
:param role: IAM role supplied for job execution. (templated)
:param rule_set_names: A list of ruleset names for evaluation. (templated)
:param number_of_workers: The number of G.1X workers to be used in the run. (default: 5)
:param timeout: The timeout for a run in minutes. This is the maximum time that a run can consume resources
before it is terminated and enters TIMEOUT status. (default: 2,880)
:param verify_result_status: Validate all the ruleset rules evaluation run results,
If any of the rule status is Fail or Error then an exception is thrown. (default: True)
:param show_results: Displays all the ruleset rules evaluation run results. (default: True)
:param rule_set_evaluation_run_kwargs: Extra arguments for evaluation run. (templated)
:param wait_for_completion: Whether to wait for job to stop. (default: True)
:param waiter_delay: Time in seconds to wait between status checks. (default: 60)
:param waiter_max_attempts: Maximum number of attempts to check for job completion. (default: 20)
:param deferrable: If True, the operator will wait asynchronously for the job to stop.
This implies waiting for completion. This mode requires aiobotocore module to be installed.
(default: False)
:param aws_conn_id: The Airflow connection used for AWS credentials.
If this is ``None`` or empty then the default boto3 behaviour is used. If
running Airflow in a distributed manner and aws_conn_id is None or
empty, then default boto3 configuration would be used (and must be
maintained on each worker node).
:param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
:param verify: Whether or not to verify SSL certificates. See:
https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
:param botocore_config: Configuration dictionary (key-values) for botocore client. See:
https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
"""
[docs]
aws_hook_class = GlueDataQualityHook
[docs]
template_fields: Sequence[str] = aws_template_fields(
"datasource",
"role",
"rule_set_names",
"rule_set_evaluation_run_kwargs",
)
[docs]
template_fields_renderers = {"datasource": "json", "rule_set_evaluation_run_kwargs": "json"}
def __init__(
self,
*,
datasource: dict,
role: str,
rule_set_names: list[str],
number_of_workers: int = 5,
timeout: int = 2880,
verify_result_status: bool = True,
show_results: bool = True,
rule_set_evaluation_run_kwargs: dict[str, Any] | None = None,
wait_for_completion: bool = True,
waiter_delay: int = 60,
waiter_max_attempts: int = 20,
deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
**kwargs,
):
super().__init__(**kwargs)
[docs]
self.datasource = datasource
[docs]
self.rule_set_names = rule_set_names
[docs]
self.number_of_workers = number_of_workers
[docs]
self.verify_result_status = verify_result_status
[docs]
self.show_results = show_results
[docs]
self.rule_set_evaluation_run_kwargs = rule_set_evaluation_run_kwargs or {}
[docs]
self.wait_for_completion = wait_for_completion
[docs]
self.waiter_delay = waiter_delay
[docs]
self.waiter_max_attempts = waiter_max_attempts
[docs]
self.deferrable = deferrable
[docs]
def execute(self, context: Context) -> str:
self.validate_inputs()
self.log.info(
"Submitting AWS Glue data quality ruleset evaluation run for RulesetNames %s", self.rule_set_names
)
response = self.hook.conn.start_data_quality_ruleset_evaluation_run(
DataSource=self.datasource,
Role=self.role,
NumberOfWorkers=self.number_of_workers,
Timeout=self.timeout,
RulesetNames=self.rule_set_names,
**self.rule_set_evaluation_run_kwargs,
)
evaluation_run_id = response["RunId"]
message_description = (
f"AWS Glue data quality ruleset evaluation run RunId: {evaluation_run_id} to complete."
)
if self.deferrable:
self.log.info("Deferring %s", message_description)
self.defer(
trigger=GlueDataQualityRuleSetEvaluationRunCompleteTrigger(
evaluation_run_id=response["RunId"],
waiter_delay=self.waiter_delay,
waiter_max_attempts=self.waiter_max_attempts,
aws_conn_id=self.aws_conn_id,
),
method_name="execute_complete",
)
elif self.wait_for_completion:
self.log.info("Waiting for %s", message_description)
self.hook.get_waiter("data_quality_ruleset_evaluation_run_complete").wait(
RunId=evaluation_run_id,
WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
)
self.log.info(
"AWS Glue data quality ruleset evaluation run completed RunId: %s", evaluation_run_id
)
self.hook.validate_evaluation_run_results(
evaluation_run_id=evaluation_run_id,
show_results=self.show_results,
verify_result_status=self.verify_result_status,
)
else:
self.log.info("AWS Glue data quality ruleset evaluation run runId: %s.", evaluation_run_id)
return evaluation_run_id
[docs]
def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> str:
validated_event = validate_execute_complete_event(event)
if validated_event["status"] != "success":
raise AirflowException(f"Error: AWS Glue data quality ruleset evaluation run: {validated_event}")
self.hook.validate_evaluation_run_results(
evaluation_run_id=validated_event["evaluation_run_id"],
show_results=self.show_results,
verify_result_status=self.verify_result_status,
)
return validated_event["evaluation_run_id"]
[docs]
class GlueDataQualityRuleRecommendationRunOperator(AwsBaseOperator[GlueDataQualityHook]):
"""
Starts a recommendation run that is used to generate rules, Glue Data Quality analyzes the data and comes up with recommendations for a potential ruleset.
Recommendation runs are automatically deleted after 90 days.
.. seealso::
For more information on how to use this operator, take a look at the guide:
:ref:`howto/operator:GlueDataQualityRuleRecommendationRunOperator`
:param datasource: The data source (Glue table) associated with this run. (templated)
:param role: IAM role supplied for job execution. (templated)
:param number_of_workers: The number of G.1X workers to be used in the run. (default: 5)
:param timeout: The timeout for a run in minutes. This is the maximum time that a run can consume resources
before it is terminated and enters TIMEOUT status. (default: 2,880)
:param show_results: Displays the recommended ruleset (a set of rules), when recommendation run completes. (default: True)
:param recommendation_run_kwargs: Extra arguments for recommendation run. (templated)
:param wait_for_completion: Whether to wait for job to stop. (default: True)
:param waiter_delay: Time in seconds to wait between status checks. (default: 60)
:param waiter_max_attempts: Maximum number of attempts to check for job completion. (default: 20)
:param deferrable: If True, the operator will wait asynchronously for the job to stop.
This implies waiting for completion. This mode requires aiobotocore module to be installed.
(default: False)
:param aws_conn_id: The Airflow connection used for AWS credentials.
If this is ``None`` or empty then the default boto3 behaviour is used. If
running Airflow in a distributed manner and aws_conn_id is None or
empty, then default boto3 configuration would be used (and must be
maintained on each worker node).
:param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
:param verify: Whether or not to verify SSL certificates. See:
https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
:param botocore_config: Configuration dictionary (key-values) for botocore client. See:
https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
"""
[docs]
aws_hook_class = GlueDataQualityHook
[docs]
template_fields: Sequence[str] = aws_template_fields(
"datasource",
"role",
"recommendation_run_kwargs",
)
[docs]
template_fields_renderers = {"datasource": "json", "recommendation_run_kwargs": "json"}
def __init__(
self,
*,
datasource: dict,
role: str,
number_of_workers: int = 5,
timeout: int = 2880,
show_results: bool = True,
recommendation_run_kwargs: dict[str, Any] | None = None,
wait_for_completion: bool = True,
waiter_delay: int = 60,
waiter_max_attempts: int = 20,
deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
**kwargs,
):
super().__init__(**kwargs)
[docs]
self.datasource = datasource
[docs]
self.number_of_workers = number_of_workers
[docs]
self.show_results = show_results
[docs]
self.recommendation_run_kwargs = recommendation_run_kwargs or {}
[docs]
self.wait_for_completion = wait_for_completion
[docs]
self.waiter_delay = waiter_delay
[docs]
self.waiter_max_attempts = waiter_max_attempts
[docs]
self.deferrable = deferrable
[docs]
def execute(self, context: Context) -> str:
glue_table = self.datasource.get("GlueTable", {})
if not glue_table.get("DatabaseName") or not glue_table.get("TableName"):
raise AttributeError("DataSource glue table must have DatabaseName and TableName")
self.log.info("Submitting AWS Glue data quality recommendation run with %s", self.datasource)
try:
response = self.hook.conn.start_data_quality_rule_recommendation_run(
DataSource=self.datasource,
Role=self.role,
NumberOfWorkers=self.number_of_workers,
Timeout=self.timeout,
**self.recommendation_run_kwargs,
)
except ClientError as error:
raise AirflowException(
f"AWS Glue data quality recommendation run failed: {error.response['Error']['Message']}"
)
recommendation_run_id = response["RunId"]
message_description = (
f"AWS Glue data quality recommendation run RunId: {recommendation_run_id} to complete."
)
if self.deferrable:
self.log.info("Deferring %s", message_description)
self.defer(
trigger=GlueDataQualityRuleRecommendationRunCompleteTrigger(
recommendation_run_id=recommendation_run_id,
waiter_delay=self.waiter_delay,
waiter_max_attempts=self.waiter_max_attempts,
aws_conn_id=self.aws_conn_id,
),
method_name="execute_complete",
)
elif self.wait_for_completion:
self.log.info("Waiting for %s", message_description)
self.hook.get_waiter("data_quality_rule_recommendation_run_complete").wait(
RunId=recommendation_run_id,
WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
)
self.log.info(
"AWS Glue data quality recommendation run completed RunId: %s", recommendation_run_id
)
if self.show_results:
self.hook.log_recommendation_results(run_id=recommendation_run_id)
else:
self.log.info(message_description)
return recommendation_run_id
[docs]
def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> str:
validated_event = validate_execute_complete_event(event)
if validated_event["status"] != "success":
raise AirflowException(f"Error: AWS Glue data quality rule recommendation run: {validated_event}")
if self.show_results:
self.hook.log_recommendation_results(run_id=validated_event["recommendation_run_id"])
return validated_event["recommendation_run_id"]