Add glue session integrations (hook and operators) #38830

mathiaHT · 2024-04-08T11:34:03Z

GlueCreateSessionOperator and GlueDeleteSessionOperator to handle glue session in a dag. I tried to stick as much as possible to the GlueJobOperator implementation

When using dbt-glue, we need to handle when and how the session is created in order to customize it for specific jobs or task group

Taragolis · 2024-04-08T12:03:09Z

airflow/providers/amazon/aws/hooks/glue_session.py

+        else:
+            self.num_of_dpus = num_of_dpus
+
+        kwargs["client_type"] = "glue"


This looks like forth hook around glue client:

airflow.providers.amazon.aws.hooks.glue.GlueJobHook

airflow.providers.amazon.aws.hooks.glue_catalog.GlueCatalogHook

airflow.providers.amazon.aws.hooks.glue_crawler.GlueCrawlerHook

We have set a strange precedent here, if they're using the same client then the same hook class could cover them all. Or at least put them all in the same module?

But also we have the precedent now, so I'm not opposed to just following it again here.

Taragolis · 2024-04-08T12:04:34Z

airflow/providers/amazon/aws/hooks/glue_session.py

+        session_id: str | None = None,
+        desc: str | None = None,
+        iam_role_name: str | None = None,
+        iam_role_arn: str | None = None,
+        num_of_dpus: int | None = None,
+        create_session_kwargs: dict | None = None,
+        session_poll_interval: int | float = 6,
+        *args,
+        **kwargs,


Are this parameters required into the constructor? And could be this hook are thin wrapper? And provide and validate parameters into the specific hook methods?

Taragolis · 2024-04-08T12:08:07Z

airflow/providers/amazon/aws/hooks/glue_session.py

+        """Process Glue Session state while polling; used by both sync and async methods."""
+        failed_states = ["FAILED", "STOPPED", "TIMEOUT"]
+        ready_state = "READY"
+
+        if state is ready_state:
+            self.log.info("Session %s State: %s", session_id, state)
+            return {"SessionState": state}
+        if state in failed_states:
+            session_error_message = f"Exiting Session {session_id} State: {state}"
+            self.log.info(session_error_message)
+            raise AirflowException(session_error_message)
+        else:
+            self.log.info(
+                "Polling for AWS Glue Session %s current run state with status %s",
+                session_id,
+                state,
+            )
+            return None


Could this implemented as a botocore waiter?

Taragolis · 2024-04-08T12:12:14Z

airflow/providers/amazon/aws/operators/glue_session.py

+    from airflow.utils.context import Context
+
+
+class GlueCreateSessionOperator(BaseOperator):


Please consider to use airflow.providers.amazon.aws.operators.base_aws.AwsBaseOperator

Some examples:

https://github.com/apache/airflow/blob/main/airflow/providers/amazon/aws/operators/appflow.py

https://github.com/apache/airflow/blob/main/airflow/providers/amazon/aws/operators/athena.py

https://github.com/apache/airflow/blob/main/airflow/providers/amazon/aws/operators/bedrock.py

Issue for the reference: #35278

Taragolis · 2024-04-08T12:12:29Z

airflow/providers/amazon/aws/operators/glue_session.py

+        "iam_role_name",
+        "iam_role_arn",
+    )
+    template_ext: Sequence[str] = ()


Suggested change

template_ext: Sequence[str] = ()

Taragolis · 2024-04-08T12:12:39Z

airflow/providers/amazon/aws/operators/glue_session.py

+
+    operator_extra_links = ()


Suggested change

operator_extra_links = ()

Taragolis · 2024-04-08T12:14:08Z

airflow/providers/amazon/aws/operators/glue_session.py

+        self.deferrable = deferrable
+
+    @cached_property
+    def glue_session_hook(self) -> GlueSessionHook:


The common name for the hooks in AWS operators/sensors is hook

Suggested change

def glue_session_hook(self) -> GlueSessionHook:

def hook(self) -> GlueSessionHook:

If you use AwsBaseOperator this will define automatically

Taragolis · 2024-04-08T12:16:16Z

airflow/providers/amazon/aws/triggers/glue_session.py

+    def __init__(
+        self,
+        session_id: str,
+        aws_conn_id: str | None,


There is also common hook settings exists: region_name, verify and botocore_config ?

Taragolis · 2024-04-08T12:27:10Z

airflow/providers/amazon/aws/triggers/glue_session.py

+        hook = GlueSessionHook(aws_conn_id=self.aws_conn_id, session_poll_interval=self.session_poll_interval)
+        await hook.async_session_readiness(self.session_id)
+        yield TriggerEvent({"status": "ready", "message": "Session ready", "value": self.session_id})


I guess in case of the error this trigger do not yield anything

o-nikolas

Would you be interested in adding/updating system tests for this? You can see inspiration in tests/system/providers/amazon/aws/

o-nikolas · 2024-04-16T23:59:40Z

airflow/providers/amazon/aws/operators/glue_session.py

+    def __init__(
+        self,
+        *,
+        session_id: str = "aws_glue_default_session",


I don't know much about glue sessions, should this be a unique uuid generated for each invocation of the operator?

o-nikolas · 2024-04-17T00:00:39Z

airflow/providers/amazon/aws/hooks/glue_session.py

+        else:
+            self.num_of_dpus = num_of_dpus
+
+        kwargs["client_type"] = "glue"


We have set a strange precedent here, if they're using the same client then the same hook class could cover them all. Or at least put them all in the same module?

But also we have the precedent now, so I'm not opposed to just following it again here.

o-nikolas · 2024-04-17T00:02:18Z

airflow/providers/amazon/aws/operators/glue_session.py

+        elif worker_type_exists and not num_workers_exists:
+            raise ValueError("Need to specify NumberOfWorkers when specifying custom WorkerType")
+        elif num_of_dpus is None:
+            self.num_of_dpus: int | float = 10


Is this a common default? If not it might be work mentioning it in the :param: doc string for num_of_dpus

This is the default value of the resource in aws

o-nikolas · 2024-04-17T00:04:16Z

airflow/providers/amazon/aws/operators/glue_session.py

+        iam_role_arn: str | None = None,
+        num_of_dpus: int | None = None,
+        create_session_kwargs: dict | None = None,
+        wait_for_readiness: bool = True,


We usually call this wait_for_completion

Also can you document this in the doc string?

eladkal · 2024-07-06T10:11:47Z

@mathiaHT are you still working on this PR?

mathiaHT · 2024-07-06T11:37:42Z

Hello, yes I will need these features at some point. But I currently have other priorities

github-actions · 2024-09-09T00:15:24Z

This pull request has been automatically marked as stale because it has not had recent activity. It will be closed in 5 days if no further activity occurs. Thank you for your contributions.

mathiaHT requested review from eladkal and o-nikolas as code owners April 8, 2024 11:34

boring-cyborg bot added area:providers provider:amazon AWS/Amazon - related issues labels Apr 8, 2024

mathiaHT changed the title ~~create glue session provider operators~~ create aws provider operators for glue session Apr 8, 2024

Taragolis reviewed Apr 8, 2024

View reviewed changes

mathiaHT requested review from potiuk, ashb and jedcunningham as code owners April 16, 2024 08:23

mathiaHT force-pushed the feat-glue-interactive-session-operators branch from 4620399 to 787a8cf Compare April 16, 2024 08:24

o-nikolas reviewed Apr 17, 2024

View reviewed changes

eladkal changed the title ~~create aws provider operators for glue session~~ Add glue session integrations (hook and operators) Apr 24, 2024

Mathia Haure-Touze and others added 3 commits May 14, 2024 10:12

create glue session provider operators

287d562

rework glue_session resources

ab8f817

update generate breeze doc

54d4298

mathiaHT force-pushed the feat-glue-interactive-session-operators branch from 787a8cf to 54d4298 Compare May 14, 2024 10:47

Mathia Haure-Touze added 2 commits May 14, 2024 15:25

fix assert called with args not kwargs

f68511e

add dag example

294e37d

wip

34edce9

github-actions bot added the stale Stale PRs per the .github/workflows/stale.yml policy file label Sep 9, 2024

github-actions bot closed this Sep 22, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add glue session integrations (hook and operators) #38830

Add glue session integrations (hook and operators) #38830

mathiaHT commented Apr 8, 2024 •

edited

Loading

Taragolis Apr 8, 2024

o-nikolas Apr 17, 2024

Taragolis Apr 8, 2024

Taragolis Apr 8, 2024

Taragolis Apr 8, 2024

Taragolis Apr 8, 2024

Taragolis Apr 8, 2024

Taragolis Apr 8, 2024

Taragolis Apr 8, 2024

Taragolis Apr 8, 2024

o-nikolas left a comment

o-nikolas Apr 16, 2024

o-nikolas Apr 17, 2024

o-nikolas Apr 17, 2024

mathiaHT May 14, 2024

o-nikolas Apr 17, 2024

eladkal commented Jul 6, 2024

mathiaHT commented Jul 6, 2024

github-actions bot commented Sep 9, 2024

		from airflow.utils.context import Context


		class GlueCreateSessionOperator(BaseOperator):

	def glue_session_hook(self) -> GlueSessionHook:
	def hook(self) -> GlueSessionHook:

Add glue session integrations (hook and operators) #38830

Add glue session integrations (hook and operators) #38830

Conversation

mathiaHT commented Apr 8, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

o-nikolas left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

eladkal commented Jul 6, 2024

mathiaHT commented Jul 6, 2024

github-actions bot commented Sep 9, 2024

mathiaHT commented Apr 8, 2024 •

edited

Loading