content: #  Copyright (c) ZenML GmbH 2022. All Rights Reserved.
  #
  #  Licensed 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:
  #
  #       https://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.
  """Module to generate an Airflow DAG from a config file."""
  
  import datetime
  import importlib
  import os
  import zipfile
  from typing import Any, Dict, List, Optional, Type, Union
  
  from pydantic import BaseModel, Field
  
  ENV_ZENML_AIRFLOW_RUN_ID = "ZENML_AIRFLOW_RUN_ID"
  ENV_ZENML_LOCAL_STORES_PATH = "ZENML_LOCAL_STORES_PATH"
  CONFIG_FILENAME = "config.json"
  
  
  class TaskConfiguration(BaseModel):
      """Airflow task configuration."""
  
      id: str
      zenml_step_name: str
      upstream_steps: List[str]
  
      docker_image: str
      command: List[str]
      arguments: List[str]
  
      environment: Dict[str, str] = {}
  
      operator_source: str
      operator_args: Dict[str, Any] = {}
  
  
  class DagConfiguration(BaseModel):
      """Airflow DAG configuration."""
  
      id: str
      tasks: List[TaskConfiguration]
  
      local_stores_path: Optional[str] = None
  
      schedule: Union[datetime.timedelta, str] = Field(
          union_mode="left_to_right"
      )
      start_date: datetime.datetime
      end_date: Optional[datetime.datetime] = None
      catchup: bool = False
  
      tags: List[str] = []
      dag_args: Dict[str, Any] = {}
  
  
  def import_class_by_path(class_path: str) -> Type[Any]:
      """Imports a class based on a given path.
  
      Args:
          class_path: str, class_source e.g. this.module.Class
  
      Returns:
          the given class
      """
      module_name, class_name = class_path.rsplit(".", 1)
      module = importlib.import_module(module_name)
      return getattr(module, class_name)  # type: ignore[no-any-return]
  
  
  def get_operator_init_kwargs(
      operator_class: Type[Any],
      dag_config: DagConfiguration,
      task_config: TaskConfiguration,
  ) -> Dict[str, Any]:
      """Gets keyword arguments to pass to the operator init method.
  
      Args:
          operator_class: The operator class for which to get the kwargs.
          dag_config: The configuration of the DAG.
          task_config: The configuration of the task.
  
      Returns:
          The init keyword arguments.
      """
      init_kwargs = {"task_id": task_config.id}
  
      try:
          from airflow.providers.docker.operators.docker import DockerOperator
  
          if issubclass(operator_class, DockerOperator):
              init_kwargs.update(
                  get_docker_operator_init_kwargs(
                      dag_config=dag_config, task_config=task_config
                  )
              )
      except ImportError:
          pass
  
      try:
          from airflow.providers.cncf.kubernetes.operators.kubernetes_pod import (
              KubernetesPodOperator,
          )
  
          if issubclass(operator_class, KubernetesPodOperator):
              init_kwargs.update(
                  get_kubernetes_pod_operator_init_kwargs(
                      dag_config=dag_config, task_config=task_config
                  )
              )
      except ImportError:
          pass
  
      try:
          # Support for apache-airflow-providers-cncf-kubernetes>=10.0.0 where
          # the import changed
          from airflow.providers.cncf.kubernetes.operators.pod import (
              KubernetesPodOperator,
          )
  
          if issubclass(operator_class, KubernetesPodOperator):
              init_kwargs.update(
                  get_kubernetes_pod_operator_init_kwargs(
                      dag_config=dag_config, task_config=task_config
                  )
              )
      except ImportError:
          pass
  
      init_kwargs.update(task_config.operator_args)
      return init_kwargs
  
  
  def get_docker_operator_init_kwargs(
      dag_config: DagConfiguration, task_config: TaskConfiguration
  ) -> Dict[str, Any]:
      """Gets keyword arguments to pass to the DockerOperator.
  
      Args:
          dag_config: The configuration of the DAG.
          task_config: The configuration of the task.
  
      Returns:
          The init keyword arguments.
      """
      mounts = []
      extra_hosts = {}
      environment = task_config.environment
      environment[ENV_ZENML_AIRFLOW_RUN_ID] = "{{run_id}}"
  
      if dag_config.local_stores_path:
          from docker.types import Mount
  
          environment[ENV_ZENML_LOCAL_STORES_PATH] = dag_config.local_stores_path
          mounts = [
              Mount(
                  target=dag_config.local_stores_path,
                  source=dag_config.local_stores_path,
                  type="bind",
              )
          ]
          extra_hosts = {"host.docker.internal": "host-gateway"}
      return {
          "image": task_config.docker_image,
          "command": task_config.command + task_config.arguments,
          "mounts": mounts,
          "environment": environment,
          "extra_hosts": extra_hosts,
      }
  
  
  def get_kubernetes_pod_operator_init_kwargs(
      dag_config: DagConfiguration, task_config: TaskConfiguration
  ) -> Dict[str, Any]:
      """Gets keyword arguments to pass to the KubernetesPodOperator.
  
      Args:
          dag_config: The configuration of the DAG.
          task_config: The configuration of the task.
  
      Returns:
          The init keyword arguments.
      """
      from kubernetes.client.models import V1EnvVar
  
      environment = task_config.environment
      environment[ENV_ZENML_AIRFLOW_RUN_ID] = "{{run_id}}"
  
      return {
          "name": f"{dag_config.id}_{task_config.id}",
          "namespace": "default",
          "image": task_config.docker_image,
          "cmds": task_config.command,
          "arguments": task_config.arguments,
          "env_vars": [
              V1EnvVar(name=key, value=value)
              for key, value in environment.items()
          ],
      }
  
  
  try:
      archive = zipfile.ZipFile(os.path.dirname(__file__), "r")
  except (IsADirectoryError, PermissionError):
      # Not inside a zip, this happens if we import this file outside of an
      # airflow dag zip
      pass
  else:
      import airflow
  
      config_str = archive.read(CONFIG_FILENAME)
      dag_config = DagConfiguration.model_validate_json(config_str)
  
      step_name_to_airflow_operator = {}
  
      with airflow.DAG(
          dag_id=dag_config.id,
          is_paused_upon_creation=False,
          tags=dag_config.tags,
          schedule=dag_config.schedule,
          start_date=dag_config.start_date,
          end_date=dag_config.end_date,
          catchup=dag_config.catchup,
          **dag_config.dag_args,
      ) as dag:
          for task in dag_config.tasks:
              operator_class = import_class_by_path(task.operator_source)
              init_kwargs = get_operator_init_kwargs(
                  operator_class=operator_class,
                  dag_config=dag_config,
                  task_config=task,
              )
              operator = operator_class(**init_kwargs)
  
              step_name_to_airflow_operator[task.zenml_step_name] = operator
              for upstream_step in task.upstream_steps:
                  operator.set_upstream(
                      step_name_to_airflow_operator[upstream_step]
                  )
  