def try_get_airflow_context():
# type: ()-> Optional[AirflowTaskContext]
# first try to get from spark
try:
from_spark = try_get_airflow_context_from_spark_conf()
if from_spark:
return from_spark
else:
_debug_init_print("couldn't get airflow context from spark")
# Those env vars are set by airflow before running the operator
dag_id = os.environ.get("AIRFLOW_CTX_DAG_ID")
execution_date = os.environ.get("AIRFLOW_CTX_EXECUTION_DATE")
task_id = os.environ.get("AIRFLOW_CTX_TASK_ID")
try:
try_number = _get_try_number()
except Exception:
try_number = None
if dag_id and task_id and execution_date:
return AirflowTaskContext(
dag_id=dag_id,
execution_date=execution_date,
task_id=task_id,
try_number=try_number,
)
return None
except Exception:
return None
def verify_spark_pre_conditions():
if spark_tracking_enabled() and _SPARK_ENV_FLAG in os.environ:
if _is_dbnd_spark_installed():
return True
else:
_debug_init_print("failed to import pyspark or dbnd-spark")
else:
_debug_init_print(
"DBND__ENABLE__SPARK_CONTEXT_ENV or SPARK_ENV_LOADED are not set")
return False
def read_spark_environ_config():
logger.debug("running read_spark_environ_config")
from pyspark import SparkContext
_debug_init_print("creating spark context to read spark conf")
spark_conf = SparkContext.getOrCreate().getConf()
spark_conf = dict(spark_conf.getAll())
spark_conf = {key.lstrip("spark.env."): value for key, value in spark_conf.items()}
return get_environ_config_from_dict(spark_conf, "environ")
def dbnd_setup_plugin():
from dbnd_spark.local.local_spark_config import SparkLocalEngineConfig
from dbnd_spark.spark_bootstrap import dbnd_spark_bootstrap
register_config_cls(SparkLocalEngineConfig)
register_config_cls(LivySparkConfig)
dbnd_spark_bootstrap()
if has_pyspark_imported() and spark_tracking_enabled():
config_store = read_spark_environ_config()
dbnd_config.set_values(config_store, "system")
else:
_debug_init_print(
"spark conf is not loaded since pyspark is not imported or DBND__ENABLE__SPARK_CONTEXT_ENV is not set"
)
def try_get_airflow_context_from_spark_conf():
# type: ()-> Optional[AirflowTaskContext]
if not spark_tracking_enabled() or _SPARK_ENV_FLAG not in os.environ:
_debug_init_print(
"DBND__ENABLE__SPARK_CONTEXT_ENV or SPARK_ENV_LOADED are not set")
return None
if not _is_dbnd_spark_installed():
_debug_init_print("failed to import pyspark or dbnd-spark")
return None
try:
_debug_init_print("creating spark context to get spark conf")
from pyspark import SparkContext
conf = SparkContext.getOrCreate().getConf()
dag_id = conf.get("spark.env.AIRFLOW_CTX_DAG_ID")
execution_date = conf.get("spark.env.AIRFLOW_CTX_EXECUTION_DATE")
task_id = conf.get("spark.env.AIRFLOW_CTX_TASK_ID")
try_number = conf.get("spark.env.AIRFLOW_CTX_TRY_NUMBER")
if dag_id and task_id and execution_date:
return AirflowTaskContext(
dag_id=dag_id,
execution_date=execution_date,
task_id=task_id,
try_number=try_number,
)
except Exception as ex:
logger.info("Failed to get airflow context info from spark job: %s",
ex)
return None
def dbnd_spark_bootstrap():
global _DBND_REGISTER_SPARK_TYPES
if _DBND_REGISTER_SPARK_TYPES:
return
# don't run it twice or in recursion
_DBND_REGISTER_SPARK_TYPES = True
_workaround_spark_namedtuple_serialization()
try:
_debug_init_print("importing pyspark")
import pyspark
except import_errors as ex:
# pyspark is not installed, user will not be able to use pyspark types
return
# we register spark types only if we have spark installed
try:
from dbnd_spark.spark_targets import dbnd_register_spark_types
dbnd_register_spark_types()
except Exception:
pass
return