def run_once(
ctx,
output: Optional[str],
dag_name: Optional[str],
wait_for_completion: Optional[int],
image: Optional[str],
): # pylint: disable=too-many-arguments
"""
Uploads pipeline to Airflow and runs once
"""
dag_filename, template_stream = get_dag_filename_and_template_stream(
ctx, dag_name=dag_name, image=image, cron_expression=None)
context_helper = ctx.obj["context_helper"]
output = output or context_helper.config.output
with fsspec.open(f"{output}/{dag_filename}", "wt") as f:
template_stream.dump(f)
airflow_client = AirflowClient(context_helper.config.host)
dag = airflow_client.wait_for_dag(
dag_id=dag_name or context_helper.config.run_config.run_name,
tag=f'commit_sha:{context_helper.session.store["git"]["commit_sha"]}',
)
dag_run_id = airflow_client.trigger_dag_run(dag.dag_id)
if (wait_for_completion or 0) > 0:
assert (airflow_client.wait_for_dag_run_completion(
dag.dag_id, dag_run_id, wait_for_completion) == "success")
def compile(ctx, image, target_path="dags/"):
"""Create an Airflow DAG for a project"""
dag_filename, template_stream = get_dag_filename_and_template_stream(
ctx, image=image, cron_expression=get_cron_expression(ctx))
target_path = Path(target_path) / dag_filename
with fsspec.open(str(target_path), "wt") as f:
template_stream.dump(f)
def upload_pipeline(ctx, output: str, image: str):
"""
Uploads pipeline to Airflow DAG location
"""
dag_filename, template_stream = get_dag_filename_and_template_stream(
ctx, image=image, cron_expression=get_cron_expression(ctx))
output = output or ctx.obj["context_helper"].config.output
with fsspec.open(f"{output}/{dag_filename}", "wt") as f:
template_stream.dump(f)
def schedule(ctx, output: str, cron_expression: str):
"""
Uploads pipeline to Airflow with given schedule
"""
dag_filename, template_stream = get_dag_filename_and_template_stream(
ctx, cron_expression=get_cron_expression(ctx, cron_expression))
output = output or ctx.obj["context_helper"].config.output
with fsspec.open(f"{output}/{dag_filename}", "wt") as f:
template_stream.dump(f)