[英]Extend BigQueryExecuteQueryOperator with additional labels using jinja2
為了使用標簽跟蹤 GCP 成本,希望使用一些額外的標簽擴展BigQueryExecuteQueryOperator ,以便每個任務實例在其構造函數中自動設置這些標簽。
class ExtendedBigQueryExecuteQueryOperator(BigQueryExecuteQueryOperator):
@apply_defaults
def __init__(self,
*args,
**kwargs) -> None:
task_labels = {
'dag_id': '{{ dag.dag_id }}',
'task_id': kwargs.get('task_id'),
'ds': '{{ ds }}',
# ugly, all three params got in diff. ways
}
super().__init__(*args, **kwargs)
if self.labels is None:
self.labels = task_labels
else:
self.labels.update(task_labels)
with DAG(dag_id=...,
start_date=...,
schedule_interval=...,
default_args=...) as dag:
t1 = ExtendedBigQueryExecuteQueryOperator(
task_id=f't1',
sql=f'SELECT 1;',
labels={'some_additional_label2':'some_additional_label2'}
# all labels should be: dag_id, task_id, ds, some_additional_label2
)
t2 = ExtendedBigQueryExecuteQueryOperator(
task_id=f't2',
sql=f'SELECT 2;',
labels={'some_additional_label3':'some_additional_label3'}
# all labels should be: dag_id, task_id, ds, some_additional_label3
)
t1 >> t2
但后來我丟失了任務級別標簽some_additional_label2
或some_additional_label3
。
您可以在airflow_local_settings.py
中創建以下策略:
def policy(task):
if task.__class__.__name__ == "BigQueryExecuteQueryOperator":
task.labels.update({'dag_id': task.dag_id, 'task_id': task.task_id})
來自文檔:
您的本地 Airflow 設置文件可以定義策略 function,該策略能夠根據其他任務或 DAG 屬性改變任務屬性。 它接收一個參數作為對任務對象的引用,並期望改變它的屬性。
有關應用策略的更多詳細信息: https://airflow.readthedocs.io/en/1.10.9/concepts.html#cluster-policy
在這種情況下,您不需要擴展 BigQueryExecuteQueryOperator。 唯一缺少的部分是您可以在任務本身中設置的execution_date 。
例子:
with DAG(dag_id=...,
start_date=...,
schedule_interval=...,
default_args=...) as dag:
t1 = BigQueryExecuteQueryOperator(
task_id=f't1',
sql=f'SELECT 1;',
lables={'some_additional_label2':'some_additional_label2', 'ds': '{{ ds }}'}
)
airflow_local_settings
文件需要在您的PYTHONPATH上。 您可以放在$AIRFLOW_HOME/config
下或 dags 目錄中。
聲明:本站的技術帖子網頁,遵循CC BY-SA 4.0協議,如果您需要轉載,請注明本站網址或者原文地址。任何問題請咨詢:yoyou2525@163.com.