AIP-72: Allow retrieving Connection from Task Context

[kaxil]

part of #44481

• Added a minimal Connection user-facing object in Task SDK definition for use in the DAG file
• Added logic to get Connections in the context. Fixed some bugs in the way related to Connection parsing/serializing!

Now, we have following Connection related objects:

• ConnectionResponse is auto-generated and tightly coupled with the API schema.
• ConnectionResult is runtime-specific and meant for internal communication between Supervisor & Task Runner.
• Connection class here is where the public-facing, user-relevant aspects are exposed, hiding internal details.

Next up:

• Same for XCom & Variable
• Implementation of BaseHook.get_conn

Tested it with a DAG:

DAG:

from __future__ importannotationsfromairflow.models.baseoperatorimportBaseOperatorfromairflow.models.dagimportdagclassCustomOperator(BaseOperator):
defexecute(self, context):
importosos.environ["AIRFLOW_CONN_AIRFLOW_DB"] ="sqlite:///home/airflow/airflow.db"task_id=context["task_instance"].task_idprint(f"Hello World {task_id}!")
print(context)
print(context["conn"].airflow_db)
assertcontext["conn"].airflow_db.conn_id=="airflow_db"@dag()defsuper_basic_run():
CustomOperator(task_id="hello")
super_basic_run()

For case where a connection is not found

---

^ Add meaningful description above
Read the Pull Request Guidelines for more information.
In case of fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in a newsfragment file, named {pr_number}.significant.rst or {issue_number}.significant.rst, in newsfragments.

[ashb]

My comments have been addressed

[amoghrajesh]

Some comments lying open, rest LGTM

[kaxil]

part of #44481

• Added a minimal Connection user-facing object in Task SDK definition for use in the DAG file
• Added logic to get Connections in the context. Fixed some bugs in the way related to Connection parsing/serializing!

Now, we have following Connection related objects:

• ConnectionResponse is auto-generated and tightly coupled with the API schema.
• ConnectionResult is runtime-specific and meant for internal communication between Supervisor & Task Runner.
• Connection class here is where the public-facing, user-relevant aspects are exposed, hiding internal details.

Next up:

• Same for XCom & Variable
• Implementation of BaseHook.get_conn

Tested it with a DAG:

DAG:

from __future__ importannotationsfromairflow.models.baseoperatorimportBaseOperatorfromairflow.models.dagimportdagclassCustomOperator(BaseOperator):
defexecute(self, context):
importosos.environ["AIRFLOW_CONN_AIRFLOW_DB"] ="sqlite:///home/airflow/airflow.db"task_id=context["task_instance"].task_idprint(f"Hello World {task_id}!")
print(context)
print(context["conn"].airflow_db)
assertcontext["conn"].airflow_db.conn_id=="airflow_db"@dag()defsuper_basic_run():
CustomOperator(task_id="hello")
super_basic_run()

For case where a connection is not found

---

^ Add meaningful description above
Read the Pull Request Guidelines for more information.
In case of fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in a newsfragment file, named {pr_number}.significant.rst or {issue_number}.significant.rst, in newsfragments.

[ashb]

My comments have been addressed

[amoghrajesh]

Some comments lying open, rest LGTM

[kaxil]

part of #44481

• Added a minimal Connection user-facing object in Task SDK definition for use in the DAG file
• Added logic to get Connections in the context. Fixed some bugs in the way related to Connection parsing/serializing!

Now, we have following Connection related objects:

• ConnectionResponse is auto-generated and tightly coupled with the API schema.
• ConnectionResult is runtime-specific and meant for internal communication between Supervisor & Task Runner.
• Connection class here is where the public-facing, user-relevant aspects are exposed, hiding internal details.

Next up:

• Same for XCom & Variable
• Implementation of BaseHook.get_conn

Tested it with a DAG:

DAG:

from __future__ importannotationsfromairflow.models.baseoperatorimportBaseOperatorfromairflow.models.dagimportdagclassCustomOperator(BaseOperator):
defexecute(self, context):
importosos.environ["AIRFLOW_CONN_AIRFLOW_DB"] ="sqlite:///home/airflow/airflow.db"task_id=context["task_instance"].task_idprint(f"Hello World {task_id}!")
print(context)
print(context["conn"].airflow_db)
assertcontext["conn"].airflow_db.conn_id=="airflow_db"@dag()defsuper_basic_run():
CustomOperator(task_id="hello")
super_basic_run()

For case where a connection is not found

---

^ Add meaningful description above
Read the Pull Request Guidelines for more information.
In case of fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in a newsfragment file, named {pr_number}.significant.rst or {issue_number}.significant.rst, in newsfragments.

[ashb]

My comments have been addressed

[amoghrajesh]

Some comments lying open, rest LGTM

[kaxil]

part of #44481

• Added a minimal Connection user-facing object in Task SDK definition for use in the DAG file
• Added logic to get Connections in the context. Fixed some bugs in the way related to Connection parsing/serializing!

Now, we have following Connection related objects:

• ConnectionResponse is auto-generated and tightly coupled with the API schema.
• ConnectionResult is runtime-specific and meant for internal communication between Supervisor & Task Runner.
• Connection class here is where the public-facing, user-relevant aspects are exposed, hiding internal details.

Next up:

• Same for XCom & Variable
• Implementation of BaseHook.get_conn

Tested it with a DAG:

DAG:

from __future__ importannotationsfromairflow.models.baseoperatorimportBaseOperatorfromairflow.models.dagimportdagclassCustomOperator(BaseOperator):
defexecute(self, context):
importosos.environ["AIRFLOW_CONN_AIRFLOW_DB"] ="sqlite:///home/airflow/airflow.db"task_id=context["task_instance"].task_idprint(f"Hello World {task_id}!")
print(context)
print(context["conn"].airflow_db)
assertcontext["conn"].airflow_db.conn_id=="airflow_db"@dag()defsuper_basic_run():
CustomOperator(task_id="hello")
super_basic_run()

For case where a connection is not found

---

^ Add meaningful description above
Read the Pull Request Guidelines for more information.
In case of fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in a newsfragment file, named {pr_number}.significant.rst or {issue_number}.significant.rst, in newsfragments.

[ashb]

My comments have been addressed

[amoghrajesh]

Some comments lying open, rest LGTM

[kaxil]

part of #44481

• Added a minimal Connection user-facing object in Task SDK definition for use in the DAG file
• Added logic to get Connections in the context. Fixed some bugs in the way related to Connection parsing/serializing!

Now, we have following Connection related objects:

• ConnectionResponse is auto-generated and tightly coupled with the API schema.
• ConnectionResult is runtime-specific and meant for internal communication between Supervisor & Task Runner.
• Connection class here is where the public-facing, user-relevant aspects are exposed, hiding internal details.

Next up:

• Same for XCom & Variable
• Implementation of BaseHook.get_conn

Tested it with a DAG:

DAG:

from __future__ importannotationsfromairflow.models.baseoperatorimportBaseOperatorfromairflow.models.dagimportdagclassCustomOperator(BaseOperator):
defexecute(self, context):
importosos.environ["AIRFLOW_CONN_AIRFLOW_DB"] ="sqlite:///home/airflow/airflow.db"task_id=context["task_instance"].task_idprint(f"Hello World {task_id}!")
print(context)
print(context["conn"].airflow_db)
assertcontext["conn"].airflow_db.conn_id=="airflow_db"@dag()defsuper_basic_run():
CustomOperator(task_id="hello")
super_basic_run()

For case where a connection is not found

---

^ Add meaningful description above
Read the Pull Request Guidelines for more information.
In case of fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in a newsfragment file, named {pr_number}.significant.rst or {issue_number}.significant.rst, in newsfragments.

[ashb]

My comments have been addressed

[amoghrajesh]

Some comments lying open, rest LGTM

[kaxil]

part of #44481

• Added a minimal Connection user-facing object in Task SDK definition for use in the DAG file
• Added logic to get Connections in the context. Fixed some bugs in the way related to Connection parsing/serializing!

Now, we have following Connection related objects:

• ConnectionResponse is auto-generated and tightly coupled with the API schema.
• ConnectionResult is runtime-specific and meant for internal communication between Supervisor & Task Runner.
• Connection class here is where the public-facing, user-relevant aspects are exposed, hiding internal details.

Next up:

• Same for XCom & Variable
• Implementation of BaseHook.get_conn

Tested it with a DAG:

DAG:

from __future__ importannotationsfromairflow.models.baseoperatorimportBaseOperatorfromairflow.models.dagimportdagclassCustomOperator(BaseOperator):
defexecute(self, context):
importosos.environ["AIRFLOW_CONN_AIRFLOW_DB"] ="sqlite:///home/airflow/airflow.db"task_id=context["task_instance"].task_idprint(f"Hello World {task_id}!")
print(context)
print(context["conn"].airflow_db)
assertcontext["conn"].airflow_db.conn_id=="airflow_db"@dag()defsuper_basic_run():
CustomOperator(task_id="hello")
super_basic_run()

For case where a connection is not found

---

^ Add meaningful description above
Read the Pull Request Guidelines for more information.
In case of fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in a newsfragment file, named {pr_number}.significant.rst or {issue_number}.significant.rst, in newsfragments.

[ashb]

My comments have been addressed

[amoghrajesh]

Some comments lying open, rest LGTM

[kaxil]

part of #44481

• Added a minimal Connection user-facing object in Task SDK definition for use in the DAG file
• Added logic to get Connections in the context. Fixed some bugs in the way related to Connection parsing/serializing!

Now, we have following Connection related objects:

• ConnectionResponse is auto-generated and tightly coupled with the API schema.
• ConnectionResult is runtime-specific and meant for internal communication between Supervisor & Task Runner.
• Connection class here is where the public-facing, user-relevant aspects are exposed, hiding internal details.

Next up:

• Same for XCom & Variable
• Implementation of BaseHook.get_conn

Tested it with a DAG:

DAG:

from __future__ importannotationsfromairflow.models.baseoperatorimportBaseOperatorfromairflow.models.dagimportdagclassCustomOperator(BaseOperator):
defexecute(self, context):
importosos.environ["AIRFLOW_CONN_AIRFLOW_DB"] ="sqlite:///home/airflow/airflow.db"task_id=context["task_instance"].task_idprint(f"Hello World {task_id}!")
print(context)
print(context["conn"].airflow_db)
assertcontext["conn"].airflow_db.conn_id=="airflow_db"@dag()defsuper_basic_run():
CustomOperator(task_id="hello")
super_basic_run()

For case where a connection is not found

---

^ Add meaningful description above
Read the Pull Request Guidelines for more information.
In case of fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in a newsfragment file, named {pr_number}.significant.rst or {issue_number}.significant.rst, in newsfragments.

[ashb]

My comments have been addressed

[amoghrajesh]

Some comments lying open, rest LGTM

[kaxil]

part of #44481

• Added a minimal Connection user-facing object in Task SDK definition for use in the DAG file
• Added logic to get Connections in the context. Fixed some bugs in the way related to Connection parsing/serializing!

Now, we have following Connection related objects:

• ConnectionResponse is auto-generated and tightly coupled with the API schema.
• ConnectionResult is runtime-specific and meant for internal communication between Supervisor & Task Runner.
• Connection class here is where the public-facing, user-relevant aspects are exposed, hiding internal details.

Next up:

• Same for XCom & Variable
• Implementation of BaseHook.get_conn

Tested it with a DAG:

DAG:

from __future__ importannotationsfromairflow.models.baseoperatorimportBaseOperatorfromairflow.models.dagimportdagclassCustomOperator(BaseOperator):
defexecute(self, context):
importosos.environ["AIRFLOW_CONN_AIRFLOW_DB"] ="sqlite:///home/airflow/airflow.db"task_id=context["task_instance"].task_idprint(f"Hello World {task_id}!")
print(context)
print(context["conn"].airflow_db)
assertcontext["conn"].airflow_db.conn_id=="airflow_db"@dag()defsuper_basic_run():
CustomOperator(task_id="hello")
super_basic_run()

For case where a connection is not found

---

^ Add meaningful description above
Read the Pull Request Guidelines for more information.
In case of fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in a newsfragment file, named {pr_number}.significant.rst or {issue_number}.significant.rst, in newsfragments.

[ashb]

My comments have been addressed

[amoghrajesh]

Some comments lying open, rest LGTM