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 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 """ 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)} """ 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 copy_helper_py_file(target_file_path: str) -> None: """ Copies. airflow/utils/helper.py to a new location within provider package The helper has two methods (chain, cross_downstream) that are moved from the original helper to 'airflow.models.baseoperator'. so in 1.10 they should reimport the original 'airflow.utils.helper' methods. Those deprecated methods use import with import_string("<IMPORT>") so it is easier to replace them as strings rather than with Bowler :param target_file_path: target path name for the helpers.py """ source_helper_file_path = os.path.join(get_source_airflow_folder(), "airflow", "utils", "helpers.py") with open(source_helper_file_path, "rt") as in_file: with open(target_file_path, "wt") as out_file: for line in in_file: out_file.write( line.replace('airflow.models.baseoperator', 'airflow.utils.helpers'))
def refactor_google_package(self): r""" 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.-]+$') """ def google_package_filter(node: LN, capture: Capture, filename: Filename) -> bool: return filename.startswith("./airflow/providers/google/") def pure_airflow_models_filter(node: LN, capture: Capture, filename: Filename) -> bool: """Check if select is exactly [airflow, . , models]""" return len(list(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"))