Add more refactor steps for providers.google (#8010)

turbaszek · web-flow · commit 8e897800716c · 2020-03-31T14:06:12.000+02:00
* Add more refactor steps for providers.google

* fixup! Add more refactor steps for providers.google
diff --git a/airflow/providers/google/cloud/example_dags/example_kubernetes_engine_setup.py b/airflow/providers/google/cloud/example_dags/example_kubernetes_engine_setup.py
diff --git a/backport_packages/setup_backport_packages.py b/backport_packages/setup_backport_packages.py
@@ -172,15 +172,7 @@ def add_provide_context_to_python_operator(node: LN, capture: Capture, filename:
 
     def remove_class(qry, class_name) -> None:
         def _remover(node: LN, capture: Capture, filename: Filename) -> None:
-            if node.type == 300:
-                for ch in node.post_order():
-                    if isinstance(ch, Leaf) and ch.value == class_name:
-                        if ch.next_sibling and ch.next_sibling.value == ",":
-                            ch.next_sibling.remove()
-                        ch.remove()
-            elif node.type == 311:
-                node.parent.remove()
-            else:
+            if node.type not in (300, 311):  # remove only definition
                 node.remove()
 
         qry.select_class(class_name).modify(_remover)
@@ -189,6 +181,10 @@ def _remover(node: LN, capture: Capture, filename: Filename) -> None:
         ("airflow.operators.bash", "airflow.operators.bash_operator"),
         ("airflow.operators.python", "airflow.operators.python_operator"),
         ("airflow.utils.session", "airflow.utils.db"),
+        (
+            "airflow.providers.cncf.kubernetes.operators.kubernetes_pod",
+            "airflow.contrib.operators.kubernetes_pod_operator"
+        ),
     ]
 
     qry = Query()
@@ -222,10 +218,20 @@ def _remover(node: LN, capture: Capture, filename: Filename) -> None:
     # Remove tags
     qry.select_method("DAG").is_call().modify(remove_tags_modifier)
 
-    # Fix KubernetesPodOperator imports to use old path
-    qry.select_module(
-        "airflow.providers.cncf.kubernetes.operators.kubernetes_pod").rename(
-        "airflow.contrib.operators.kubernetes_pod_operator"
+    # Fix AWS import in Google Cloud Transfer Service
+    (
+        qry
+        .select_module("airflow.providers.amazon.aws.hooks.base_aws")
+        .is_filename(include=r"cloud_storage_transfer_service\.py")
+        .rename("airflow.contrib.hooks.aws_hook")
+    )
+
+    (
+        qry
+        .select_class("AwsBaseHook")
+        .is_filename(include=r"cloud_storage_transfer_service\.py")
+        .filter(lambda n, c, f: n.type == 300)
+        .rename("AwsHook")
     )
 
     # Fix BaseOperatorLinks imports
@@ -243,10 +249,28 @@ def _remover(node: LN, capture: Capture, filename: Filename) -> None:
         .modify(add_provide_context_to_python_operator)
     )
 
+    # Remove new class and rename usages of old
     remove_class(qry, "GKEStartPodOperator")
+    (
+        qry
+        .select_class("GKEStartPodOperator")
+        .is_filename(include=r"example_kubernetes_engine\.py")
+        .rename("GKEPodOperator")
+    )
 
     qry.execute(write=True, silent=False, interactive=False)
 
+    # Add old import to GKE
+    gke_path = os.path.join(
+        dirname(__file__), "airflow", "providers", "google", "cloud", "operators", "kubernetes_engine.py"
+    )
+    with open(gke_path, "a") as f:
+        f.writelines(["", "from airflow.contrib.operators.gcp_container_operator import GKEPodOperator"])
+
+    gke_path = os.path.join(
+        dirname(__file__), "airflow", "providers", "google", "cloud", "operators", "kubernetes_engine.py"
+    )
+
 
 def get_source_providers_folder():
     return os.path.join(dirname(__file__), os.pardir, "airflow", "providers")
diff --git a/tests/providers/google/cloud/operators/test_kubernetes_engine_system.py b/tests/providers/google/cloud/operators/test_kubernetes_engine_system.py
@@ -24,11 +24,6 @@
 @pytest.mark.backend("mysql", "postgres")
 @pytest.mark.credential_file(GCP_GKE_KEY)
 class KubernetesEngineExampleDagTest(GoogleSystemTest):
-    @pytest.mark.airflow_2
     @provide_gcp_context(GCP_GKE_KEY)
     def test_run_example_gcp_gke(self):
         self.run_dag('example_gcp_gke', CLOUD_DAG_FOLDER)
-
-    @provide_gcp_context(GCP_GKE_KEY)
-    def test_run_example_gcp_gke_setup(self):
-        self.run_dag('example_gcp_gke_setup', CLOUD_DAG_FOLDER)
