docs/exts/providers_extensions.py - airflow - Git at Google

 # Licensed to the Apache Software Foundation (ASF) under one
 # or more contributor license agreements.  See the NOTICE file
 # distributed with this work for additional information
 # regarding copyright ownership.  The ASF licenses this file
 # to you under the Apache License, Version 2.0 (the
 # "License"); you may not use this file except in compliance
 # with the License.  You may obtain a copy of the License at
 #
 #   http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing,
 # software distributed under the License is distributed on an
 # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
 """Module for provider's custom Sphinx extensions that will be loaded in all providers' documentation."""

 from __future__ import annotations

 import ast
 import os
 from collections.abc import Iterable
 from functools import partial
 from pathlib import Path
 from typing import Any, Callable

 # No stub exists for docutils.parsers.rst.directives. See https://github.com/python/typeshed/issues/5755.
 from provider_yaml_utils import load_package_data

 from docs.exts.operators_and_hooks_ref import (
     DEFAULT_HEADER_SEPARATOR,
     BaseJinjaReferenceDirective,
     _render_template,
 )


 def find_class_methods_with_specific_calls(
     class_node: ast.ClassDef, target_calls: set[str], import_mappings: dict[str, str]
 ) -> set[str]:
     """
     Identifies class methods that make specific calls.

     This function only tracks target calls within the class scope. Method calling some function defined
     will not be taken into consideration even if this function performs a target call.

     Method calling other method that performs a target call will also be included.

     This function performs a two-pass analysis of the AST:
     1. It first identifies methods containing direct calls to the specified functions
        and records method calls on `self`.
     2. It then identifies methods that indirectly make such calls by invoking the
        methods identified in the first pass.

     :param class_node: The root node of the AST representing the class to analyze.
     :param target_calls: A set of full paths to the method names to track when called.
     :param import_mappings: A mapping of import names to fully qualified module names.

     :return: Method names within the class that either directly or indirectly make the specified calls.

     Examples:
         > source_code = '''
         ... class Example:
         ...     def method1(self):
         ...         my_method().ok()

         ...     def method2(self):
         ...         self.method1()

         ...     def method3(self):
         ...         my_method().not_ok()

         ...     def method4(self):
         ...         self.some_other_method()
         ... '''
         > find_methods_with_specific_calls(
             ast.parse(source_code),
             {"airflow.my_method.not_ok", "airflow.my_method.ok"},
             {"my_method": "airflow.my_method"}
         )
         {'method1', 'method2', 'method3'}
     """
     method_call_map: dict[str, set[str]] = {}
     methods_with_calls: set[str] = set()

     # First pass: Collect all calls and identify methods with specific calls we are looking for
     for node in ast.walk(class_node):
         if not isinstance(node, ast.FunctionDef):
             continue
         method_call_map[node.name] = set()
         for sub_node in ast.walk(node):
             if not isinstance(sub_node, ast.Call):
                 continue
             called_function = sub_node.func
             if not isinstance(called_function, ast.Attribute):
                 continue
             if isinstance(called_function.value, ast.Call) and isinstance(
                 called_function.value.func, ast.Name
             ):
                 full_method_call = (
                     f"{import_mappings.get(called_function.value.func.id)}.{called_function.attr}"
                 )
                 if full_method_call in target_calls:
                     methods_with_calls.add(node.name)
             elif isinstance(called_function.value, ast.Name) and called_function.value.id == "self":
                 method_call_map[node.name].add(called_function.attr)

     # Second pass: Identify all methods that call the ones in `methods_with_calls`
     def find_calling_methods(method_name):
         for caller, callees in method_call_map.items():
             if method_name in callees and caller not in methods_with_calls:
                 methods_with_calls.add(caller)
                 find_calling_methods(caller)

     for method in list(methods_with_calls):
         find_calling_methods(method)

     return methods_with_calls


 def get_import_mappings(tree) -> dict[str, str]:
     """Retrieve a mapping of local import names to their fully qualified module paths from an AST tree.

     :param tree: The AST tree to analyze for import statements.

     :return: A dictionary where the keys are the local names (aliases) used in the current module
         and the values are the fully qualified names of the imported modules or their members.

     Example:
         >>> import ast
         >>> code = '''
         ... import os
         ... import numpy as np
         ... from collections import defaultdict
         ... from datetime import datetime as dt
         ... '''
         >>> get_import_mappings(ast.parse(code))
         {'os': 'os', 'np': 'numpy', 'defaultdict': 'collections.defaultdict', 'dt': 'datetime.datetime'}
     """
     imports = {}
     for node in ast.walk(tree):
         if isinstance(node, (ast.Import, ast.ImportFrom)):
             for alias in node.names:
                 module_prefix = f"{node.module}." if hasattr(node, "module") and node.module else ""
                 imports[alias.asname or alias.name] = f"{module_prefix}{alias.name}"
     return imports


 def _get_module_class_registry(
     module_filepath: Path, module_name: str, class_extras: dict[str, Callable]
 ) -> dict[str, dict[str, Any]]:
     """Extracts classes and its information from a Python module file.

     The function parses the specified module file and registers all classes.
     The registry for each class includes the module filename, methods, base classes
     and any additional class extras provided.

     :param module_filepath: The file path of the module.
     :param class_extras: Additional information to include in each class's registry.

     :return: A dictionary with class names as keys and their corresponding information.
     """
     with open(module_filepath) as file:
         ast_obj = ast.parse(file.read())

     import_mappings = get_import_mappings(ast_obj)
     module_class_registry = {
         f"{module_name}.{node.name}": {
             "methods": {n.name for n in ast.walk(node) if isinstance(n, ast.FunctionDef)},
             "base_classes": [
                 import_mappings.get(b.id, f"{module_name}.{b.id}")
                 for b in node.bases
                 if isinstance(b, ast.Name)
             ],
             **{
                 key: callable_(class_node=node, import_mappings=import_mappings)
                 for key, callable_ in class_extras.items()
             },
         }
         for node in ast_obj.body
         if isinstance(node, ast.ClassDef)
     }
     return module_class_registry


 def _has_method(
     class_path: str, method_names: Iterable[str], class_registry: dict[str, dict[str, Any]]
 ) -> bool:
     """Determines if a class or its bases in the registry have any of the specified methods.

     :param class_path: The path of the class to check.
     :param method_names: A list of names of methods to search for.
     :param class_registry: A dictionary representing the class registry, where each key is a class name
                             and the value is its metadata.
     :return: True if any of the specified methods are found in the class or its base classes; False otherwise.

     Example:
     >>> example_class_registry = {
     ...     "some.module.MyClass": {"methods": {"foo", "bar"}, "base_classes": ["BaseClass"]},
     ...     "another.module.BaseClass": {"methods": {"base_foo"}, "base_classes": []},
     ... }
     >>> _has_method("some.module.MyClass", ["foo"], example_class_registry)
     True
     >>> _has_method("some.module.MyClass", ["base_foo"], example_class_registry)
     True
     >>> _has_method("some.module.MyClass", ["not_a_method"], example_class_registry)
     False
     """
     if class_path in class_registry:
         if any(method in class_registry[class_path]["methods"] for method in method_names):
             return True
         for base_name in class_registry[class_path]["base_classes"]:
             if _has_method(base_name, method_names, class_registry):
                 return True
     return False


 def _get_providers_class_registry(
     class_extras: dict[str, Callable] | None = None,
 ) -> dict[str, dict[str, Any]]:
     """Builds a registry of classes from YAML configuration files.

     This function scans through YAML configuration files to build a registry of classes.
     It parses each YAML file to get the provider's name and registers classes from Python
     module files within the provider's directory, excluding '__init__.py'.

     :return: A dictionary with provider names as keys and a dictionary of classes as values.
     """
     class_registry = {}
     for provider_yaml_content in load_package_data():
         provider_pkg_root = Path(provider_yaml_content["package-dir"])
         for root, _, file_names in os.walk(provider_pkg_root):
             folder = Path(root)
             for file_name in file_names:
                 if not file_name.endswith(".py") or file_name == "__init__.py":
                     continue

                 module_filepath = folder.joinpath(file_name)

                 module_registry = _get_module_class_registry(
                     module_filepath=module_filepath,
                     module_name=(
                         provider_yaml_content["python-module"]
                         + "."
                         + module_filepath.relative_to(provider_pkg_root)
                         .with_suffix("")
                         .as_posix()
                         .replace("/", ".")
                     ),
                     class_extras={
                         "provider_name": lambda **kwargs: provider_yaml_content["package-name"],
                         **(class_extras or {}),
                     },
                 )
                 class_registry.update(module_registry)

     return class_registry


 def _render_openlineage_supported_classes_content():
     openlineage_operator_methods = ("get_openlineage_facets_on_complete", "get_openlineage_facets_on_start")
     openlineage_db_hook_methods = (
         "get_openlineage_database_info",
         "get_openlineage_database_specific_lineage",
     )
     hook_lineage_collector_path = "airflow.providers.common.compat.lineage.hook.get_hook_lineage_collector"
     hook_level_lineage_collector_calls = {
         f"{hook_lineage_collector_path}.add_input_asset",  # Airflow 3
         f"{hook_lineage_collector_path}.add_output_asset",  # Airflow 3
         f"{hook_lineage_collector_path}.add_input_dataset",  # Airflow 2
         f"{hook_lineage_collector_path}.add_output_dataset",  # Airflow 2
     }

     class_registry = _get_providers_class_registry(
         class_extras={
             "methods_with_hook_level_lineage": partial(
                 find_class_methods_with_specific_calls, target_calls=hook_level_lineage_collector_calls
             )
         }
     )

     # These excluded classes will be included in docs directly
     class_registry.pop("airflow.providers.common.sql.hooks.sql.DbApiHook")
     class_registry.pop("airflow.providers.common.sql.operators.sql.SQLExecuteQueryOperator")

     providers: dict[str, dict[str, Any]] = {}
     db_hooks: list[tuple[str, str]] = []
     for class_path, info in class_registry.items():
         class_name = class_path.split(".")[-1]
         if class_name.startswith("_"):
             continue
         provider_entry = providers.setdefault(info["provider_name"], {"operators": {}, "hooks": {}})

         if class_name.lower().endswith("operator"):
             if _has_method(
                 class_path=class_path,
                 method_names=openlineage_operator_methods,
                 class_registry=class_registry,
             ):
                 provider_entry["operators"][class_path] = [
                     f"{class_path}.{method}"
                     for method in set(openlineage_operator_methods) & set(info["methods"])
                 ]
         elif class_name.lower().endswith("hook"):
             if _has_method(
                 class_path=class_path,
                 method_names=openlineage_db_hook_methods,
                 class_registry=class_registry,
             ):
                 db_type = class_name.replace("SqlApiHook", "").replace("Hook", "")
                 db_hooks.append((db_type, class_path))

             elif info["methods_with_hook_level_lineage"]:
                 provider_entry["hooks"][class_path] = [
                     f"{class_path}.{method}"
                     for method in info["methods_with_hook_level_lineage"]
                     if not method.startswith("_")
                 ]

     providers = {
         provider: {
             "operators": {
                 operator: sorted(methods)
                 for operator, methods in sorted(
                     details["operators"].items(), key=lambda x: x[0].split(".")[-1]
                 )
             },
             "hooks": {
                 hook: sorted(methods)
                 for hook, methods in sorted(details["hooks"].items(), key=lambda x: x[0].split(".")[-1])
             },
         }
         for provider, details in sorted(providers.items())
         if any(details.values())  # This filters out providers with empty 'operators' and 'hooks'
     }
     db_hooks = sorted({db_type: hook for db_type, hook in db_hooks}.items(), key=lambda x: x[0])

     return _render_template(
         "openlineage.rst.jinja2",
         providers=providers,
         db_hooks=db_hooks,
     )


 class OpenLineageSupportedClassesDirective(BaseJinjaReferenceDirective):
     """Generate list of classes supporting OpenLineage"""

     def render_content(self, *, tags: set[str] | None, header_separator: str = DEFAULT_HEADER_SEPARATOR):
         return _render_openlineage_supported_classes_content()


 def setup(app):
     """Setup plugin"""
     app.add_directive("airflow-providers-openlineage-supported-classes", OpenLineageSupportedClassesDirective)

     return {"parallel_read_safe": True, "parallel_write_safe": True}
	# Licensed to the Apache Software Foundation (ASF) under one
	# or more contributor license agreements. See the NOTICE file
	# distributed with this work for additional information
	# regarding copyright ownership. The ASF licenses this file
	# to you under the Apache License, Version 2.0 (the
	# "License"); you may not use this file except in compliance
	# with the License. You may obtain a copy of the License at
	#
	# http://www.apache.org/licenses/LICENSE-2.0
	#
	# Unless required by applicable law or agreed to in writing,
	# software distributed under the License is distributed on an
	# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
	# KIND, either express or implied. See the License for the
	# specific language governing permissions and limitations
	# under the License.
	"""Module for provider's custom Sphinx extensions that will be loaded in all providers' documentation."""

	from __future__ import annotations

	import ast
	import os
	from collections.abc import Iterable
	from functools import partial
	from pathlib import Path
	from typing import Any, Callable

	# No stub exists for docutils.parsers.rst.directives. See https://github.com/python/typeshed/issues/5755.
	from provider_yaml_utils import load_package_data

	from docs.exts.operators_and_hooks_ref import (
	DEFAULT_HEADER_SEPARATOR,
	BaseJinjaReferenceDirective,
	_render_template,
	)


	def find_class_methods_with_specific_calls(
	class_node: ast.ClassDef, target_calls: set[str], import_mappings: dict[str, str]
	) -> set[str]:
	"""
	Identifies class methods that make specific calls.

	This function only tracks target calls within the class scope. Method calling some function defined
	will not be taken into consideration even if this function performs a target call.

	Method calling other method that performs a target call will also be included.

	This function performs a two-pass analysis of the AST:
	1. It first identifies methods containing direct calls to the specified functions
	and records method calls on `self`.
	2. It then identifies methods that indirectly make such calls by invoking the
	methods identified in the first pass.

	:param class_node: The root node of the AST representing the class to analyze.
	:param target_calls: A set of full paths to the method names to track when called.
	:param import_mappings: A mapping of import names to fully qualified module names.

	:return: Method names within the class that either directly or indirectly make the specified calls.

	Examples:
	> source_code = '''
	... class Example:
	... def method1(self):
	... my_method().ok()

	... def method2(self):
	... self.method1()

	... def method3(self):
	... my_method().not_ok()

	... def method4(self):
	... self.some_other_method()
	... '''
	> find_methods_with_specific_calls(
	ast.parse(source_code),
	{"airflow.my_method.not_ok", "airflow.my_method.ok"},
	{"my_method": "airflow.my_method"}
	)
	{'method1', 'method2', 'method3'}
	"""
	method_call_map: dict[str, set[str]] = {}
	methods_with_calls: set[str] = set()

	# First pass: Collect all calls and identify methods with specific calls we are looking for
	for node in ast.walk(class_node):
	if not isinstance(node, ast.FunctionDef):
	continue
	method_call_map[node.name] = set()
	for sub_node in ast.walk(node):
	if not isinstance(sub_node, ast.Call):
	continue
	called_function = sub_node.func
	if not isinstance(called_function, ast.Attribute):
	continue
	if isinstance(called_function.value, ast.Call) and isinstance(
	called_function.value.func, ast.Name
	):
	full_method_call = (
	f"{import_mappings.get(called_function.value.func.id)}.{called_function.attr}"
	)
	if full_method_call in target_calls:
	methods_with_calls.add(node.name)
	elif isinstance(called_function.value, ast.Name) and called_function.value.id == "self":
	method_call_map[node.name].add(called_function.attr)

	# Second pass: Identify all methods that call the ones in `methods_with_calls`
	def find_calling_methods(method_name):
	for caller, callees in method_call_map.items():
	if method_name in callees and caller not in methods_with_calls:
	methods_with_calls.add(caller)
	find_calling_methods(caller)

	for method in list(methods_with_calls):
	find_calling_methods(method)

	return methods_with_calls


	def get_import_mappings(tree) -> dict[str, str]:
	"""Retrieve a mapping of local import names to their fully qualified module paths from an AST tree.

	:param tree: The AST tree to analyze for import statements.

	:return: A dictionary where the keys are the local names (aliases) used in the current module
	and the values are the fully qualified names of the imported modules or their members.

	Example:
	>>> import ast
	>>> code = '''
	... import os
	... import numpy as np
	... from collections import defaultdict
	... from datetime import datetime as dt
	... '''
	>>> get_import_mappings(ast.parse(code))
	{'os': 'os', 'np': 'numpy', 'defaultdict': 'collections.defaultdict', 'dt': 'datetime.datetime'}
	"""
	imports = {}
	for node in ast.walk(tree):
	if isinstance(node, (ast.Import, ast.ImportFrom)):
	for alias in node.names:
	module_prefix = f"{node.module}." if hasattr(node, "module") and node.module else ""
	imports[alias.asname or alias.name] = f"{module_prefix}{alias.name}"
	return imports


	def _get_module_class_registry(
	module_filepath: Path, module_name: str, class_extras: dict[str, Callable]
	) -> dict[str, dict[str, Any]]:
	"""Extracts classes and its information from a Python module file.

	The function parses the specified module file and registers all classes.
	The registry for each class includes the module filename, methods, base classes
	and any additional class extras provided.

	:param module_filepath: The file path of the module.
	:param class_extras: Additional information to include in each class's registry.

	:return: A dictionary with class names as keys and their corresponding information.
	"""
	with open(module_filepath) as file:
	ast_obj = ast.parse(file.read())

	import_mappings = get_import_mappings(ast_obj)
	module_class_registry = {
	f"{module_name}.{node.name}": {
	"methods": {n.name for n in ast.walk(node) if isinstance(n, ast.FunctionDef)},
	"base_classes": [
	import_mappings.get(b.id, f"{module_name}.{b.id}")
	for b in node.bases
	if isinstance(b, ast.Name)
	],
	**{
	key: callable_(class_node=node, import_mappings=import_mappings)
	for key, callable_ in class_extras.items()
	},
	}
	for node in ast_obj.body
	if isinstance(node, ast.ClassDef)
	}
	return module_class_registry


	def _has_method(
	class_path: str, method_names: Iterable[str], class_registry: dict[str, dict[str, Any]]
	) -> bool:
	"""Determines if a class or its bases in the registry have any of the specified methods.

	:param class_path: The path of the class to check.
	:param method_names: A list of names of methods to search for.
	:param class_registry: A dictionary representing the class registry, where each key is a class name
	and the value is its metadata.
	:return: True if any of the specified methods are found in the class or its base classes; False otherwise.

	Example:
	>>> example_class_registry = {
	... "some.module.MyClass": {"methods": {"foo", "bar"}, "base_classes": ["BaseClass"]},
	... "another.module.BaseClass": {"methods": {"base_foo"}, "base_classes": []},
	... }
	>>> _has_method("some.module.MyClass", ["foo"], example_class_registry)
	True
	>>> _has_method("some.module.MyClass", ["base_foo"], example_class_registry)
	True
	>>> _has_method("some.module.MyClass", ["not_a_method"], example_class_registry)
	False
	"""
	if class_path in class_registry:
	if any(method in class_registry[class_path]["methods"] for method in method_names):
	return True
	for base_name in class_registry[class_path]["base_classes"]:
	if _has_method(base_name, method_names, class_registry):
	return True
	return False


	def _get_providers_class_registry(
	class_extras: dict[str, Callable] \| None = None,
	) -> dict[str, dict[str, Any]]:
	"""Builds a registry of classes from YAML configuration files.

	This function scans through YAML configuration files to build a registry of classes.
	It parses each YAML file to get the provider's name and registers classes from Python
	module files within the provider's directory, excluding '__init__.py'.

	:return: A dictionary with provider names as keys and a dictionary of classes as values.
	"""
	class_registry = {}
	for provider_yaml_content in load_package_data():
	provider_pkg_root = Path(provider_yaml_content["package-dir"])
	for root, _, file_names in os.walk(provider_pkg_root):
	folder = Path(root)
	for file_name in file_names:
	if not file_name.endswith(".py") or file_name == "__init__.py":
	continue

	module_filepath = folder.joinpath(file_name)

	module_registry = _get_module_class_registry(
	module_filepath=module_filepath,
	module_name=(
	provider_yaml_content["python-module"]
	+ "."
	+ module_filepath.relative_to(provider_pkg_root)
	.with_suffix("")
	.as_posix()
	.replace("/", ".")
	),
	class_extras={
	"provider_name": lambda **kwargs: provider_yaml_content["package-name"],
	**(class_extras or {}),
	},
	)
	class_registry.update(module_registry)

	return class_registry


	def _render_openlineage_supported_classes_content():
	openlineage_operator_methods = ("get_openlineage_facets_on_complete", "get_openlineage_facets_on_start")
	openlineage_db_hook_methods = (
	"get_openlineage_database_info",
	"get_openlineage_database_specific_lineage",
	)
	hook_lineage_collector_path = "airflow.providers.common.compat.lineage.hook.get_hook_lineage_collector"
	hook_level_lineage_collector_calls = {
	f"{hook_lineage_collector_path}.add_input_asset", # Airflow 3
	f"{hook_lineage_collector_path}.add_output_asset", # Airflow 3
	f"{hook_lineage_collector_path}.add_input_dataset", # Airflow 2
	f"{hook_lineage_collector_path}.add_output_dataset", # Airflow 2
	}

	class_registry = _get_providers_class_registry(
	class_extras={
	"methods_with_hook_level_lineage": partial(
	find_class_methods_with_specific_calls, target_calls=hook_level_lineage_collector_calls
	)
	}
	)

	# These excluded classes will be included in docs directly
	class_registry.pop("airflow.providers.common.sql.hooks.sql.DbApiHook")
	class_registry.pop("airflow.providers.common.sql.operators.sql.SQLExecuteQueryOperator")

	providers: dict[str, dict[str, Any]] = {}
	db_hooks: list[tuple[str, str]] = []
	for class_path, info in class_registry.items():
	class_name = class_path.split(".")[-1]
	if class_name.startswith("_"):
	continue
	provider_entry = providers.setdefault(info["provider_name"], {"operators": {}, "hooks": {}})

	if class_name.lower().endswith("operator"):
	if _has_method(
	class_path=class_path,
	method_names=openlineage_operator_methods,
	class_registry=class_registry,
	):
	provider_entry["operators"][class_path] = [
	f"{class_path}.{method}"
	for method in set(openlineage_operator_methods) & set(info["methods"])
	]
	elif class_name.lower().endswith("hook"):
	if _has_method(
	class_path=class_path,
	method_names=openlineage_db_hook_methods,
	class_registry=class_registry,
	):
	db_type = class_name.replace("SqlApiHook", "").replace("Hook", "")
	db_hooks.append((db_type, class_path))

	elif info["methods_with_hook_level_lineage"]:
	provider_entry["hooks"][class_path] = [
	f"{class_path}.{method}"
	for method in info["methods_with_hook_level_lineage"]
	if not method.startswith("_")
	]

	providers = {
	provider: {
	"operators": {
	operator: sorted(methods)
	for operator, methods in sorted(
	details["operators"].items(), key=lambda x: x[0].split(".")[-1]
	)
	},
	"hooks": {
	hook: sorted(methods)
	for hook, methods in sorted(details["hooks"].items(), key=lambda x: x[0].split(".")[-1])
	},
	}
	for provider, details in sorted(providers.items())
	if any(details.values()) # This filters out providers with empty 'operators' and 'hooks'
	}
	db_hooks = sorted({db_type: hook for db_type, hook in db_hooks}.items(), key=lambda x: x[0])

	return _render_template(
	"openlineage.rst.jinja2",
	providers=providers,
	db_hooks=db_hooks,
	)


	class OpenLineageSupportedClassesDirective(BaseJinjaReferenceDirective):
	"""Generate list of classes supporting OpenLineage"""

	def render_content(self, *, tags: set[str] \| None, header_separator: str = DEFAULT_HEADER_SEPARATOR):
	return _render_openlineage_supported_classes_content()


	def setup(app):
	"""Setup plugin"""
	app.add_directive("airflow-providers-openlineage-supported-classes", OpenLineageSupportedClassesDirective)

	return {"parallel_read_safe": True, "parallel_write_safe": True}