def refactor_amazon_package(self): """ Fixes to "amazon" providers package. Copies some of the classes used from core Airflow to "common.utils" package of the the provider and renames imports to use them from there. We copy typing_compat.py and change import as in example diff: .. code-block:: diff --- ./airflow/providers/amazon/aws/operators/ecs.py +++ ./airflow/providers/amazon/aws/operators/ecs.py @@ -24,7 +24,7 @@ from airflow.models import BaseOperator from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook from airflow.providers.amazon.aws.hooks.logs import AwsLogsHook -from airflow.typing_compat import Protocol, runtime_checkable +from airflow.providers.amazon.common.utils.typing_compat import Protocol, runtime_checkable from airflow.utils.decorators import apply_defaults """ # noinspection PyUnusedLocal def amazon_package_filter(node: LN, capture: Capture, filename: Filename) -> bool: return filename.startswith("./airflow/providers/amazon/") os.makedirs(os.path.join(get_target_providers_package_folder("amazon"), "common", "utils"), exist_ok=True) copyfile( os.path.join(get_source_airflow_folder(), "airflow", "utils", "__init__.py"), os.path.join(get_target_providers_package_folder("amazon"), "common", "__init__.py")) copyfile( os.path.join(get_source_airflow_folder(), "airflow", "utils", "__init__.py"), os.path.join(get_target_providers_package_folder("amazon"), "common", "utils", "__init__.py")) copyfile( os.path.join(get_source_airflow_folder(), "airflow", "typing_compat.py"), os.path.join(get_target_providers_package_folder("amazon"), "common", "utils", "typing_compat.py")) (self.qry.select_module("airflow.typing_compat").filter( callback=amazon_package_filter).rename( "airflow.providers.amazon.common.utils.typing_compat")) copyfile( os.path.join(get_source_airflow_folder(), "airflow", "utils", "email.py"), os.path.join(get_target_providers_package_folder("amazon"), "common", "utils", "email.py")) (self.qry.select_module("airflow.utils.email").filter( callback=amazon_package_filter).rename( "airflow.providers.amazon.common.utils.email"))
def refactor_odbc_package(self): """ Fixes to "odbc" providers package. Copies some of the classes used from core Airflow to "common.utils" package of the the provider and renames imports to use them from there. We copy helpers.py and change import as in example diff: .. code-block:: diff --- ./airflow/providers/google/cloud/example_dags/example_datacatalog.py +++ ./airflow/providers/google/cloud/example_dags/example_datacatalog.py @@ -37,7 +37,7 @@ CloudDataCatalogUpdateTagTemplateOperator, ) from airflow.utils.dates import days_ago -from airflow.utils.helpers import chain +from airflow.providers.odbc.utils.helpers import chain default_args = {"start_date": days_ago(1)} """ # noinspection PyUnusedLocal def odbc_package_filter(node: LN, capture: Capture, filename: Filename) -> bool: return filename.startswith("./airflow/providers/odbc/") os.makedirs(os.path.join(get_target_providers_folder(), "odbc", "utils"), exist_ok=True) copyfile( os.path.join(get_source_airflow_folder(), "airflow", "utils", "__init__.py"), os.path.join(get_target_providers_package_folder("odbc"), "utils", "__init__.py")) copy_helper_py_file( os.path.join(get_target_providers_package_folder("odbc"), "utils", "helpers.py")) (self.qry.select_module("airflow.utils.helpers").filter( callback=odbc_package_filter).rename( "airflow.providers.odbc.utils.helpers"))
def refactor_google_package(self): """ Fixes to "google" providers package. Copies some of the classes used from core Airflow to "common.utils" package of the the provider and renames imports to use them from there. Note that in this case we also rename the imports in the copied files. For example we copy python_virtualenv.py, process_utils.py and change import as in example diff: .. code-block:: diff --- ./airflow/providers/google/cloud/operators/kubernetes_engine.py +++ ./airflow/providers/google/cloud/operators/kubernetes_engine.py @@ -28,11 +28,11 @@ from airflow.exceptions import AirflowException from airflow.models import BaseOperator -from airflow.providers.cncf.kubernetes.operators.kubernetes_pod import KubernetesPodOperator +from airflow.contrib.operators.kubernetes_pod_operator import KubernetesPodOperator from airflow.providers.google.cloud.hooks.kubernetes_engine import GKEHook from airflow.providers.google.common.hooks.base_google import GoogleBaseHook from airflow.utils.decorators import apply_defaults -from airflow.utils.process_utils import execute_in_subprocess, patch_environ +from airflow.providers.google.common.utils.process_utils import execute_in_subprocess And in the copied python_virtualenv.py we also change import to process_utils.py. This happens automatically and is solved by Pybowler. .. code-block:: diff --- ./airflow/providers/google/common/utils/python_virtualenv.py +++ ./airflow/providers/google/common/utils/python_virtualenv.py @@ -21,7 +21,7 @@ \"\"\" from typing import List, Optional -from airflow.utils.process_utils import execute_in_subprocess +from airflow.providers.google.common.utils.process_utils import execute_in_subprocess def _generate_virtualenv_cmd(tmp_dir: str, python_bin: str, system_site_packages: bool) We also rename Base operator links to deprecated names: .. code-block:: diff --- ./airflow/providers/google/cloud/operators/mlengine.py +++ ./airflow/providers/google/cloud/operators/mlengine.py @@ -24,7 +24,7 @@ from typing import List, Optional from airflow.exceptions import AirflowException -from airflow.models import BaseOperator, BaseOperatorLink +from airflow.models.baseoperator import BaseOperator, BaseOperatorLink from airflow.models.taskinstance import TaskInstance from airflow.providers.google.cloud.hooks.mlengine import MLEngineHook from airflow.utils.decorators import apply_defaults We remove GKEStartPodOperator (example in remove_class method) We also copy (google.common.utils) and rename imports to the helpers. .. code-block:: diff --- ./airflow/providers/google/cloud/example_dags/example_datacatalog.py +++ ./airflow/providers/google/cloud/example_dags/example_datacatalog.py @@ -37,7 +37,7 @@ CloudDataCatalogUpdateTagTemplateOperator, ) from airflow.utils.dates import days_ago -from airflow.utils.helpers import chain +from airflow.providers.google.common.utils.helpers import chain default_args = {"start_date": days_ago(1)} And also module_loading which is used by helpers .. code-block:: diff --- ./airflow/providers/google/common/utils/helpers.py +++ ./airflow/providers/google/common/utils/helpers.py @@ -26,7 +26,7 @@ from jinja2 import Template from airflow.exceptions import AirflowException -from airflow.utils.module_loading import import_string +from airflow.providers.google.common.utils.module_loading import import_string KEY_REGEX = re.compile(r'^[\\w.-]+$') """ # noinspection PyUnusedLocal def google_package_filter(node: LN, capture: Capture, filename: Filename) -> bool: return filename.startswith("./airflow/providers/google/") # noinspection PyUnusedLocal def pure_airflow_models_filter(node: LN, capture: Capture, filename: Filename) -> bool: """Check if select is exactly [airflow, . , models]""" return len([ch for ch in node.children[1].leaves()]) == 3 os.makedirs(os.path.join(get_target_providers_package_folder("google"), "common", "utils"), exist_ok=True) copyfile( os.path.join(get_source_airflow_folder(), "airflow", "utils", "__init__.py"), os.path.join(get_target_providers_package_folder("google"), "common", "utils", "__init__.py")) copyfile( os.path.join(get_source_airflow_folder(), "airflow", "utils", "python_virtualenv.py"), os.path.join(get_target_providers_package_folder("google"), "common", "utils", "python_virtualenv.py")) copy_helper_py_file( os.path.join(get_target_providers_package_folder("google"), "common", "utils", "helpers.py")) copyfile( os.path.join(get_source_airflow_folder(), "airflow", "utils", "module_loading.py"), os.path.join(get_target_providers_package_folder("google"), "common", "utils", "module_loading.py")) (self.qry.select_module("airflow.utils.python_virtualenv").filter( callback=google_package_filter).rename( "airflow.providers.google.common.utils.python_virtualenv")) copyfile( os.path.join(get_source_airflow_folder(), "airflow", "utils", "process_utils.py"), os.path.join(get_target_providers_package_folder("google"), "common", "utils", "process_utils.py")) (self.qry.select_module("airflow.utils.process_utils").filter( callback=google_package_filter).rename( "airflow.providers.google.common.utils.process_utils")) (self.qry.select_module("airflow.utils.helpers").filter( callback=google_package_filter).rename( "airflow.providers.google.common.utils.helpers")) (self.qry.select_module("airflow.utils.module_loading").filter( callback=google_package_filter).rename( "airflow.providers.google.common.utils.module_loading")) ( # Fix BaseOperatorLinks imports self.qry.select_module("airflow.models").is_filename( include=r"bigquery\.py|mlengine\.py" ).filter(callback=google_package_filter ).filter(pure_airflow_models_filter).rename( "airflow.models.baseoperator")) self.remove_class("GKEStartPodOperator") (self.qry.select_class("GKEStartPodOperator").filter( callback=google_package_filter).is_filename( include=r"example_kubernetes_engine\.py").rename( "GKEPodOperator"))