Create inventory of Trigger, Hook & Operator's using DB directly

[kaxil]

A good amount of Trigger, Hooks & Operators in the Provider are calling DB directly. Create a list of inventory of those providers, so once AIP-72 is in a good shape, the community & provider's stakeholders can lead the effort to update those.

[potiuk]

It should also be automated/verified with tests -> as described in Stage 3 of #42632 (comment)

[kaxil]

I created a very basic list:

The following providers utilize's Airflow's session object:

• Databricks
• Edge
• Elasticsearch
• FAB
• Google
• Kubernetes
• OpenLineage
• OpenSearch
• Standard

This GitHub search should show the results too:

---

Databricks

• Uses it for the Workflow plugin

Edge

• Expected for now since it is an executor.

Elasticsearch

• Uses it in the log handler. Needs to change

airflow/providers/src/airflow/providers/elasticsearch/log/es_task_handler.py

Lines 225 to 245 in 6751c96

	def _render_log_id(self, ti: TaskInstance | TaskInstanceKey, try_number: int) -> str:
	from airflow.models.taskinstance import TaskInstanceKey
	with create_session() as session:
	if isinstance(ti, TaskInstanceKey):
	ti = _ensure_ti(ti, session)
	dag_run = ti.get_dagrun(session=session)
	if USE_PER_RUN_LOG_ID:
	log_id_template = dag_run.get_log_template(session=session).elasticsearch_id
	if TYPE_CHECKING:
	assert ti.task
	try:
	dag = ti.task.dag
	except AttributeError: # ti.task is not always set.
	data_interval = (dag_run.data_interval_start, dag_run.data_interval_end)
	else:
	if TYPE_CHECKING:
	assert dag is not None
	# TODO: Task-SDK: Where should this function be?
	data_interval = dag.get_run_data_interval(dag_run) # type: ignore[attr-defined]

FAB

• For Auth, stores data in FAB related tables in Airflow DB

Google

Uses it for getting current task instance in the Triggerer for BigQuery & Dataproc. Needs to change

airflow/providers/src/airflow/providers/google/cloud/triggers/dataproc.py

Lines 246 to 255 in 6751c96

	def safe_to_cancel(self) -> bool:
	"""
	Whether it is safe to cancel the external job which is being executed by this trigger.
	This is to avoid the case that `asyncio.CancelledError` is called because the trigger itself is stopped.
	Because in those cases, we should NOT cancel the external job.
	"""
	# Database query is needed to get the latest state of the task instance.
	task_instance = self.get_task_instance() # type: ignore[call-arg]
	return task_instance.state != TaskInstanceState.DEFERRED

K8s

Uses it for:

• Rendering K8s Pod Spec YAML to show in Airflow UI for KPO & KE. (Code
• For KubernetesExecutor. Expected for now, needs to change via Convert the KubernetesExecutor to run tasks using new Task SDK supervisor code #45427
• For helper function to look up the run_id from the TI table!

OpenLineage
Uses it for checking if a TI is rescheduled or not but might be using it for more things. Needs to change!
is_ti_rescheduled_already (Code)

OpenSearch

• Same as ES and needs to change. (Code)

Standard

• Lot of changes needed here.
• ExternalTaskSensor and its Triggers (code)
• TriggerDagRunOperator and its ExtraOperatorLinks (code)
• Bash Task decorator (code)

---

There might be other usage which I haven't dug into. I only checked for presence of the following in a file:

• from airflow.utils.session
• from airflow.settings import Session

There might be other cases where a Provider code might be importing things from airflow.models to import DagRun or TaskInstance and such directly, that will need to be found out and changed too.

---

Apart from this task itself, we will need to separately change all the provider code to not import any code from airflow core and instead strictly depend on the Task SDK.

A small script to get a similar inventory:

import os
import re
import json

# Directories to scan
DIRECTORIES = ["providers", "airflow/operators"]

# Patterns to identify Airflow metadata DB access
DB_PATTERNS = [
    r"from airflow\.utils\.session",
    r"from airflow\.settings import Session",
    # r"from airflow\.models"
]

# Exclude tests directory
EXCLUDED_DIRS = ["tests"]

OUTPUT_FILE = "metadata_db_call_inventory.json"

def matches_db_pattern(line):
    """Check if a line matches any metadata DB access pattern."""
    return any(re.search(pattern, line) for pattern in DB_PATTERNS)

def scan_file(filepath):
    """Scan a single file for metadata DB access patterns."""
    file_inventory = []
    with open(filepath, "r", encoding="utf-8") as f:
        try:
            lines = f.readlines()
        except Exception as e:
            print(f"Error reading {filepath}: {e}")
            return []

        for i, line in enumerate(lines, start=1):
            if matches_db_pattern(line):
                file_inventory.append({"line": i, "content": line.strip()})
    return file_inventory

def scan_directory(directory):
    """Scan the given directory for metadata DB access patterns."""
    inventory = {}
    for root, dirs, files in os.walk(directory):
        dirs[:] = [d for d in dirs if d not in EXCLUDED_DIRS]

        for file in files:
            if file.endswith(".py"):
                filepath = os.path.join(root, file)
                file_inventory = scan_file(filepath)
                if file_inventory:
                    rel_path = os.path.relpath(filepath)
                    inventory[rel_path] = file_inventory
    return inventory

def main():
    """Main function to generate the inventory."""
    full_inventory = {}

    for directory in DIRECTORIES:
        if os.path.exists(directory):
            print(f"Scanning {directory}...")
            inventory = scan_directory(directory)
            full_inventory.update(inventory)
        else:
            print(f"Directory not found: {directory}")

    with open(OUTPUT_FILE, "w", encoding="utf-8") as f:
        json.dump(full_inventory, f, indent=4)

    print(f"Inventory saved to {OUTPUT_FILE}")

if __name__ == "__main__":
    main()

metadata_db_call_inventory.json

[kaxil]

@potiuk Would you be able to run this program some point next month?

[potiuk]

@potiuk Would you be able to run this program some point next month?

Sure

[kaxil]

Hey Jarek, did you have a chance to work on this?

[potiuk]

Hey Jarek, did you have a chance to work on this?

Not yet - I was too busy moving distributions around. But this one is definitely next