Add docs

astronomer · Mar 19, 2024 · 2b6d4c8 · 2b6d4c8
1 parent 1d33b50
commit 2b6d4c8
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 7 deletions.
diff --git a/astronomer/providers/apache/hive/example_dags/example_hive.py b/astronomer/providers/apache/hive/example_dags/example_hive.py
@@ -212,6 +212,7 @@ def add_inbound_rule_for_security_group(task_instance: Any) -> None:
 
 
 def revoke_inbound_rules(task_instance):
+    """Remove an ingress rule from security group"""
     import boto3
 
     current_docker_ip = get("https://api.ipify.org").text
@@ -222,8 +223,8 @@ def revoke_inbound_rules(task_instance):
         CidrIp=ip_range,
         FromPort=22,
         GroupId=task_instance.xcom_pull(
-                key="cluster_response_master_security_group", task_ids=["describe_created_cluster"]
-            )[0],
+            key="cluster_response_master_security_group", task_ids=["describe_created_cluster"]
+        )[0],
         IpProtocol="tcp",
     )
     logging.info("%s", response)

diff --git a/astronomer/providers/apache/livy/example_dags/example_livy.py b/astronomer/providers/apache/livy/example_dags/example_livy.py
@@ -175,6 +175,7 @@ def add_inbound_rule_for_security_group(task_instance: Any) -> None:
 
 
 def revoke_inbound_rules(task_instance):
+    """Remove an ingress rule from security group"""
     import boto3
 
     current_docker_ip = get("https://api.ipify.org").text
@@ -186,8 +187,8 @@ def revoke_inbound_rules(task_instance):
         FromPort=LIVY_OPERATOR_INGRESS_PORT,
         ToPort=LIVY_OPERATOR_INGRESS_PORT,
         GroupId=task_instance.xcom_pull(
-                key="cluster_response_master_security_group", task_ids=["describe_created_cluster"]
-            )[0],
+            key="cluster_response_master_security_group", task_ids=["describe_created_cluster"]
+        )[0],
         IpProtocol="tcp",
     )
     logging.info("%s", response)

diff --git a/astronomer/providers/sftp/example_dags/example_sftp.py b/astronomer/providers/sftp/example_dags/example_sftp.py
@@ -158,6 +158,7 @@ def add_inbound_rule_for_security_group(task_instance: "TaskInstance") -> None:
 
 
 def revoke_inbound_rules(task_instance):
+    """Remove an ingress rule from security group"""
     import boto3
 
     current_docker_ip = get("https://api.ipify.org").text
@@ -169,8 +170,8 @@ def revoke_inbound_rules(task_instance):
         FromPort=22,
         ToPort=22,
         GroupId=task_instance.xcom_pull(
-                key="cluster_response_master_security_group", task_ids=["describe_created_cluster"]
-            )[0],
+            key="cluster_response_master_security_group", task_ids=["describe_created_cluster"]
+        )[0],
         IpProtocol="tcp",
     )
     logging.info("%s", response)
@@ -299,7 +300,9 @@ def check_dag_status(**kwargs: Any) -> None:
     )
 
     revoke_inbound_rules = PythonOperator(
-        task_id="revoke_inbound_rules", trigger_rule=TriggerRule.ALL_DONE, python_callable=revoke_inbound_rules
+        task_id="revoke_inbound_rules",
+        trigger_rule=TriggerRule.ALL_DONE,
+        python_callable=revoke_inbound_rules,
     )
 
     dag_final_status = PythonOperator(