Remove direct access to db from providers

[potiuk]

Body

Related to #45839

Found 9 providers with metadata DB access patterns:

Provider: apache-airflow-providers-cncf-kubernetes

• kubernetes_helper_functions.py
  • Line:108: from airflow.settings import Session
  • Line:118: session.query(TaskInstance.run_id)
• kubernetes_executor.py
  • Line:75: from airflow.utils.session import NEW_SESSION, provide_session
  • Line:399: @provide_session

Provider: apache-airflow-providers-databricks

• databricks_workflow.py
  • Line:43: from airflow.utils.session import NEW_SESSION, provide_session
  • Line:48: from sqlalchemy.orm.session import Session
  • Line:96: @provide_session
  • Line:109: return session.query(DagRun).filter(DagRun.dag_id == dag.dag_id, DagRun.run_id == run_id).first()
  • Line:112: @provide_session
  • Line:204: @provide_session
  • Line:212: session.query(TaskInstance)

Provider: apache-airflow-providers-edge

• edge_executor_plugin.py
  • Line:44: from airflow.utils.session import NEW_SESSION, provide_session
  • Line:133: @provide_session
  • Line:151: @provide_session
• edge_executor.py
  • Line:42: from airflow.utils.session import NEW_SESSION, provide_session
  • Line:95: @provide_session
  • Line:116: @provide_session
  • Line:146: @provide_session
  • Line:179: session.query(EdgeWorkerModel)
  • Line:205: session.query(EdgeJobModel)
  • Line:247: session.query(EdgeJobModel)
  • Line:320: @provide_session
• edge_worker.py
  • Line:32: from airflow.utils.session import NEW_SESSION, provide_session
  • Line:197: @provide_session
  • Line:208: @provide_session
  • Line:217: @provide_session
  • Line:223: @provide_session
• _v2_routes.py
  • Line:42: from airflow.utils.session import NEW_SESSION, create_session, provide_session
  • Line:130: @provide_session
  • Line:144: @provide_session
  • Line:162: @provide_session
  • Line:179: @provide_session
• logs.py
  • Line:39: from airflow.utils.session import NEW_SESSION, provide_session
  • Line:45: @provide_session

Provider: apache-airflow-providers-elasticsearch

• es_task_handler.py
  • Line:52: from airflow.utils.session import create_session
  • Line:104: session.query(TaskInstance)

Provider: apache-airflow-providers-fab

• utils.py
  • Line:43: from sqlalchemy.orm.session import Session
• fab_auth_manager.py
  • Line:101: from airflow.utils.session import NEW_SESSION, create_session, provide_session
  • Line:414: @provide_session
• override.py
  • Line:544: num_sessions = session.query(user_session_model).count()
  • Line:558: for s in session.query(user_session_model):
  • Line:1270: return self.get_session.query(self.role_model).filter_by(name=name).one_or_none()
  • Line:1273: return self.get_session.query(self.role_model).all()
  • Line:1282: role = session.query(Role).filter(Role.name == role_name).first()
  • Line:1316: return self.get_session.query(self.role_model).filter_by(name=self.auth_role_public).one_or_none()
  • Line:1367: return self.get_session.query(func.count(self.user_model.id)).scalar()
  • Line:1400: self.get_session.query(self.user_model)
  • Line:1406: self.get_session.query(self.user_model)
  • Line:1415: return self.get_session.query(self.user_model).filter_by(email=email).one_or_none()
  • Line:1447: return self.get_session.query(self.user_model).all()
  • Line:1487: return self.get_session.query(self.action_model).filter_by(name=name).one_or_none()
  • Line:1521: self.get_session.query(self.permission_model)
  • Line:1548: return self.get_session.query(self.resource_model).filter_by(name=name).one_or_none()
  • Line:1590: self.get_session.query(self.permission_model)
  • Line:1602: return self.get_session.query(self.permission_model).filter_by(resource_id=resource.id).all()
  • Line:1647: self.get_session.query(self.role_model).filter(self.role_model.permissions.contains(perm)).first()
  • Line:1657: if not self.get_session.query(self.permission_model).filter_by(action=perm.action).all():

Provider: apache-airflow-providers-google

• dataproc.py
  • Line:37: from airflow.utils.session import provide_session
  • Line:41: from sqlalchemy.orm.session import Session
  • Line:120: @provide_session
  • Line:127: query = session.query(TaskInstance).filter(
  • Line:227: @provide_session
  • Line:229: query = session.query(TaskInstance).filter(
• bigquery.py
  • Line:30: from airflow.utils.session import provide_session
  • Line:34: from sqlalchemy.orm.session import Session
  • Line:100: @provide_session
  • Line:102: query = session.query(TaskInstance).filter(

Provider: apache-airflow-providers-openlineage

• listener.py
  • Line:57: from airflow.settings import Session
• utils.py
  • Line:57: from airflow.utils.session import NEW_SESSION, provide_session
  • Line:204: @provide_session
  • Line:215: session.query(
  • Line:221: session.query(

Provider: apache-airflow-providers-opensearch

• os_task_handler.py
  • Line:43: from airflow.utils.session import create_session
  • Line:89: session.query(TaskInstance)

Provider: apache-airflow-providers-standard

• trigger_dagrun.py
  • Line:44: from airflow.utils.session import NEW_SESSION, provide_session
  • Line:53: from sqlalchemy.orm.session import Session
  • Line:340: @provide_session
• external_task.py
  • Line:30: from airflow.utils.session import NEW_SESSION, provide_session
  • Line:270: @provide_session
  • Line:279: session.query(func.count("*")) # .count() is inefficient
• sensor_helper.py
  • Line:25: from airflow.utils.session import NEW_SESSION, provide_session
  • Line:32: @provide_session
• external_task.py
  • Line:35: from airflow.utils.session import NEW_SESSION, provide_session
  • Line:403: @provide_session

Committer

• I acknowledge that I am a maintainer/committer of the Apache Airflow project.

[potiuk]

Following #45389 - here are remaining db access method in our providers (cc: @kaxil ). I marked the right people (I think that have beeen involved there - I would love to get confirmation if my assessment of those cases is correct - i will exclude / mark them as done if so:

• cncf.kubernetes -> the only access is from executor @ashb was the last one accessing it. I understand that we are still going to add cncf.provider to scheduler in order to use k8S.

• databricks -> the only access is in the workflow plugin @sunank200 @pankajkoti I think you were the authors - and this is the "legacy" plugin that is going to still have access to the database (when the provider is installed with webserver)

• edge -> it's a bit strange to see edge provider accessing the db in worker part, so that one needs to be looked at @AutomationDev85 @jscheffl

• fab -> this is the auth_manager only that should legitimately access db from the webserver. Correct @vincbeck ?

• google -> here it looks like dataproc and biguery operators are accessing the DB directly to get access to the task instance - I believe both should be changed to use the Interal API of task-sdk in this case @VladaZakharova

• openlineage > the database access in listener are also something that runs on scheduler only ? or is it something that should be fixed as well @mobuchowski @kacpermuda

• opensearch and elasticsearch are both using db when running task handler - are those task handlers already converted to use task-sdk logging or is this something we do not plan to do? @ashb @kaxil @amoghrajesh

• ``standard` - here we still have a few operators accessing db directly but I believe that is work in-progress still @amoghrajesh @kaxil @gopidesupavan

[ashb]

Executors run in the scheduler - those can still access db.

Standard provider only tries to access db on airflow 2

[pankajkoti]

The analysis is correct for the Databricks provider and it needs to be resolved while working on #40587. Please let me know if you're looking for something more there

[potiuk]

The analysis is correct for the Databricks provider and it needs to be resolved while working on #40587. Please let me know if you're looking for something more there

Cool. Thanks!

[jscheffl]

@potiuk The DB access in edge is 100% as expected.

• The edge_executor_plugin.py plugin is loaded in the webserver - atm not functioning in 3.0 as the plugin interface needs to be established. I assume will rework this then if plugin interface appears in 3.1 and need to make it a React component. Then the DB access will move to a FastAPI. Access to DB is okay in webserver though.
• The edge_executor.py - as ash mentioned already - runs in the scheduler and needs DB access to persist tasks in the job table. This is fully intended and ok.
• edge_worker.py is the shared DB model for ORM which is shared between webserver and scheduler. Direct DB access is just made there for migration (until tooling is harmonized and 2.10 support is dropped - then it is a pure ORM object). Name might be confusing but actually it is running in webserver/scheduler and the name refers to the DB table. It is NOT executed on the worker. So no reason having remote DB access.
• _v2_routes.py and logs.py is belonging to the API endpoints (hosted in webserver/api server) for the backend of the worker communication. These classes need DB access of course. Logging might need a rework in the future if the logging in TaskSDK evolves, was created before TaskSDK had anything to bring the logs from edge to cloud.

[gopidesupavan]

Yes @potiuk standard provider db access only for AF 2. :)

[amoghrajesh]

My update:

@potiuk yes the standard provider part is AF2 specific: https://github.com/apache/airflow/blob/main/providers/standard/src/airflow/providers/standard/operators/trigger_dagrun.py#L44

I think opensearch and elasticsearch is yet to be done. I'll let @ashb / @kaxil confirm.

[mobuchowski]

For OpenLineage, the code path is only for AF2 or running on scheduler. Though if top-level import is an issue, added PR to import session only on Airflow 2: #49209

[ashb]

OS and ES are fine -- the db is only used in the _read path to render the log fname

[eladkal]

Did we list somewhere how users with custom operators can remove direct access while keeping the functionality (via API)?

[potiuk]

Did we list somewhere how users with custom operators can remove direct access while keeping the functionality (via API)?

I don't think so. The direct access were never "officially supported", but I guess some guidelines should be needed

[ashb]

Did we list somewhere how users with custom operators can remove direct access while keeping the functionality (via API)?

It's too open eneded a question, as it depends entirely on what the custom operator is doing.