ExternalPython script.py
ActivePublic
Actions

Authored by atsuko on Thu, Sep 24, 12:48 PM.

Project Tags

None

Referenced Files

	F104585453: ExternalPython script.py
	Thu, Sep 24, 12:48 PM

Subscribers

None

	from __future__ import annotations

	import cloudpickle
	import sys
	import os
	# Setting the PYTHON_OPERATORS_VIRTUAL_ENV_MODE environment variable to 1,
	# helps to avoid the issue of re creating the orm session in the settings file, otherwise
	# it fails with airflow-db-not-allowed

	os.environ["PYTHON_OPERATORS_VIRTUAL_ENV_MODE"] = "1"


	if sys.version_info >= (3,6):
	try:
	from airflow.plugins_manager import integrate_macros_plugins
	integrate_macros_plugins()
	except ImportError:

	pass


	try:
	from airflow.sdk.execution_time import task_runner
	except ModuleNotFoundError:
	pass
	else:
	reinit_supervisor_comms = getattr(task_runner, "reinit_supervisor_comms", None)
	if reinit_supervisor_comms:
	reinit_supervisor_comms()

	# Script
	def _run_cleanup_task():
	import logging
	import time

	logger = logging.getLogger()
	logger.error("started cleanup")
	logger.error("sleeping 60s")
	time.sleep(60)

	retention_days = 10
	delete_limit = 10
	"""Wrapper function for PythonOperator to run cleanup with direct delete."""
	import pendulum
	from wmf_airflow_common.tasks.db_cleanup import run_cleanup_with_direct_delete

	logger.error("sleeping 60s")
	time.sleep(60)

	logger.error("started cleanup")
	data_interval_start = pendulum.today() # FIXME
	clean_before = data_interval_start.subtract(days=retention_days)
	run_cleanup_with_direct_delete(clean_before_timestamp=clean_before, delete_limit=delete_limit, dry_run=False)
	logger.error("finished cleanup, cooldown")

	logger.error("sleeping 180s")
	time.sleep(180)
	logger.error("done")
	return True


	# monkey patching for the cases when python_callable is part of the dag module.

	arg_dict = {"args": [], "kwargs": {}}




	# Read string args
	with open(sys.argv[3], "r") as file:
	virtualenv_string_args = list(map(lambda x: x.strip(), list(file)))


	try:
	res = _run_cleanup_task(arg_dict["args"], *arg_dict["kwargs"])
	except Exception as e:
	with open(sys.argv[4], "w") as file:
	file.write(str(e))
	raise

	# Write output
	with open(sys.argv[2], "wb") as file:
	if res is not None:

	cloudpickle.dump(res, file)

Event Timeline

atsuko created this paste.Thu, Sep 24, 12:48 PM

atsuko mentioned this in T434896: Fix direct database access in Airflow 3.

ExternalPython script.pyActivePublicActions

Event Timeline

ExternalPython script.py
ActivePublic
Actions