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

* Add more refactor steps for providers.google * fixup! Add more refactor steps for providers.google
apache · Mar 31, 2020 · 8e89780 · 8e89780
1 parent 07de602
commit 8e89780
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 77 deletions.
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)