SQLA2: Address assorted mypy violations in airflow-core (#59746)

* Remove redefinition of `get` overloads in parser.py

Lines 285-289 have `@overload` declarations but no implementation, and line 980 re-defines `get` with overloads.

* Add type hint for `upgraded_values`

* Fix mypy violations related to OTEL

* Add type hints to datadog's `DogStatsd` kwargs

* Address mypy violations related to structlog + refactoring

* Address mypy violations in `scheduler_job_runner`

Author: Dev-iL

airflow-core/src/airflow/jobs/scheduler_job_runner.py

@@ -2188,14 +2188,19 @@ def _schedule_dag_run(
                 ):
                     dag_model.calculate_dagrun_date_fields(dag, get_run_data_interval(dag.timetable, dag_run))
 
-                dag_run = session.scalar(
+                dag_run_reloaded = session.scalar(
                     select(DagRun)
                     .where(DagRun.id == dag_run.id)
                     .options(
                         selectinload(DagRun.consumed_asset_events).selectinload(AssetEvent.asset),
                         selectinload(DagRun.consumed_asset_events).selectinload(AssetEvent.source_aliases),
                     )
                 )
+                if dag_run_reloaded is None:
+                    # This should never happen since we just had the dag_run
+                    self.log.error("DagRun %s was deleted unexpectedly", dag_run.id)
+                    return None
+                dag_run = dag_run_reloaded
                 callback_to_execute = DagCallbackRequest(
                     filepath=dag_model.relative_fileloc or "",
                     dag_id=dag.dag_id,

shared/configuration/src/airflow_shared/configuration/parser.py

@@ -248,7 +248,7 @@ def __init__(
         self.configuration_description = configuration_description
         self._default_values = _default_values
         self._suppress_future_warnings = False
-        self.upgraded_values = {}
+        self.upgraded_values: dict[tuple[str, str], str] = {}
 
     @functools.cached_property
     def inversed_deprecated_options(self):
@@ -282,12 +282,6 @@ def sensitive_config_values(self) -> set[tuple[str, str]]:
         sensitive.update(depr_section, depr_option)
         return sensitive
 
-    @overload  # type: ignore[override]
-    def get(self, section: str, key: str, fallback: str = ..., **kwargs) -> str: ...
-
-    @overload
-    def get(self, section: str, key: str, **kwargs) -> str | None: ...
-
     def _update_defaults_from_string(self, config_string: str) -> None:
         """
         Update the defaults in _default_values based on values in config_string ("ini" format).

shared/logging/src/airflow_shared/logging/structlog.py

@@ -27,6 +27,7 @@
 from collections.abc import Callable, Iterable, Mapping, Sequence
 from functools import cache, cached_property, partial
 from pathlib import Path
+from types import ModuleType
 from typing import TYPE_CHECKING, Any, BinaryIO, Generic, TextIO, TypeVar, cast
 
 import pygtrie
@@ -96,7 +97,7 @@ def meth(self: Any, event: str, *args: Any, **kw: Any) -> Any:
 
             # See https://github.com/python/cpython/blob/3.13/Lib/logging/__init__.py#L307-L326 for reason
             if args and len(args) == 1 and isinstance(args[0], Mapping) and args[0]:
-                args = args[0]
+                return self._proxy_to_logger(name, event % args[0], **kw)
             return self._proxy_to_logger(name, event % args, **kw)
 
         meth.__name__ = name
@@ -185,14 +186,14 @@ def __init__(self, name: str | None = None, file: TextIO | None = None):
 class LoggerFactory(Generic[LogOutputType]):
     def __init__(
         self,
-        cls: Callable[[str | None, LogOutputType | None], WrappedLogger],
+        cls: type[WrappedLogger],
         io: LogOutputType | None = None,
     ):
         self.cls = cls
         self.io = io
 
     def __call__(self, logger_name: str | None = None, *args: Any) -> WrappedLogger:
-        return self.cls(logger_name, self.io)
+        return self.cls(logger_name, self.io)  # type: ignore[call-arg]
 
 
 def logger_name(logger: Any, method_name: Any, event_dict: EventDict) -> EventDict:
@@ -282,17 +283,17 @@ def structlog_processors(
 
     import click
 
-    suppress = (click, contextlib)
+    suppress: tuple[ModuleType, ...] = (click, contextlib)
     try:
         import httpcore
 
-        suppress += (httpcore,)
+        suppress = (*suppress, httpcore)
     except ImportError:
         pass
     try:
         import httpx
 
-        suppress += (httpx,)
+        suppress = (*suppress, httpx)
     except ImportError:
         pass
 
@@ -320,7 +321,8 @@ def json_dumps(msg, default):
         json = structlog.processors.JSONRenderer(serializer=json_dumps)
 
         def json_processor(logger: Any, method_name: Any, event_dict: EventDict) -> str:
-            return json(logger, method_name, event_dict).decode("utf-8")
+            result = json(logger, method_name, event_dict)
+            return result.decode("utf-8") if isinstance(result, bytes) else result
 
         shared_processors.extend(
             (
@@ -331,6 +333,7 @@ def json_processor(logger: Any, method_name: Any, event_dict: EventDict) -> str:
 
         return shared_processors, json_processor, json
 
+    exc_formatter: structlog.dev.RichTracebackFormatter | structlog.typing.ExceptionRenderer
     if os.getenv("DEV", "") != "":
         # Only use Rich in dev -- otherwise for "production" deployments it makes the logs harder to read as
         # it uses lots of ANSI escapes and non ASCII characters. Simpler is better for non-dev non-JSON
@@ -349,6 +352,7 @@ def json_processor(logger: Any, method_name: Any, event_dict: EventDict) -> str:
     if colors:
         my_styles["debug"] = structlog.dev.CYAN
 
+    console: PercentFormatRender | structlog.dev.ConsoleRenderer
     if log_format:
         console = PercentFormatRender(
             fmt=log_format,
@@ -479,15 +483,18 @@ def is_atty():
 
     wrapper_class = cast("type[BindableLogger]", make_filtering_logger())
     if json_output:
-        logger_factory = LoggerFactory(NamedBytesLogger, io=output)
+        logger_factory: LoggerFactory[Any] = LoggerFactory(NamedBytesLogger, io=output)
     else:
         # There is no universal way of telling if a file-like-object is binary (and needs bytes) or text that
         # works for files, sockets and io.StringIO/BytesIO.
 
         # If given a binary object, wrap it in a text mode wrapper
+        text_output: TextIO | None = None
         if output is not None and not hasattr(output, "encoding"):
-            output = io.TextIOWrapper(output, line_buffering=True)
-        logger_factory = LoggerFactory(NamedWriteLogger, io=output)
+            text_output = io.TextIOWrapper(cast("BinaryIO", output), line_buffering=True)
+        elif output is not None:
+            text_output = cast("TextIO", output)
+        logger_factory = LoggerFactory(NamedWriteLogger, io=text_output)
 
     structlog.configure(
         processors=shared_pre_chain + [for_structlog],

shared/observability/src/airflow_shared/observability/metrics/datadog_logger.py

@@ -175,12 +175,16 @@ def get_dogstatsd_logger(
     """Get DataDog StatsD logger."""
     from datadog import DogStatsd
 
-    dogstatsd = DogStatsd(
-        host,
-        port,
-        namespace,
-        constant_tags=cls.get_constant_tags(),
-    )
+    dogstatsd_kwargs: dict[str, str | int | list[str]] = {
+        "constant_tags": cls.get_constant_tags(),
+    }
+    if host is not None:
+        dogstatsd_kwargs["host"] = host
+    if port is not None:
+        dogstatsd_kwargs["port"] = port
+    if namespace is not None:
+        dogstatsd_kwargs["namespace"] = namespace
+    dogstatsd = DogStatsd(**dogstatsd_kwargs)
     metric_tags_validator = PatternBlockListValidator(statsd_disabled_tags)
     validator = get_validator(metrics_allow_list, metrics_block_list)
     return SafeDogStatsdLogger(

shared/observability/src/airflow_shared/observability/metrics/otel_logger.py

@@ -388,7 +388,9 @@ def get_otel_logger(
     stat_name_handler: Callable[[str], str] | None = None,
     statsd_influxdb_enabled: bool = False,
 ) -> SafeOtelLogger:
-    resource = Resource.create(attributes={SERVICE_NAME: service_name})
+    effective_service_name: str = service_name or "airflow"
+    effective_prefix: str = prefix or DEFAULT_METRIC_NAME_PREFIX
+    resource = Resource.create(attributes={SERVICE_NAME: effective_service_name})
     protocol = "https" if ssl_active else "http"
     # Allow transparent support for standard OpenTelemetry SDK environment variables.
     # https://opentelemetry.io/docs/specs/otel/protocol/exporter/#configuration-options
@@ -420,5 +422,5 @@ def get_otel_logger(
     validator = get_validator(metrics_allow_list, metrics_block_list)
 
     return SafeOtelLogger(
-        metrics.get_meter_provider(), prefix, validator, stat_name_handler, statsd_influxdb_enabled
+        metrics.get_meter_provider(), effective_prefix, validator, stat_name_handler, statsd_influxdb_enabled
     )

shared/observability/src/airflow_shared/observability/traces/otel_tracer.py

@@ -77,7 +77,7 @@ def __init__(
             log.info("(otel_tracer.__init__) - [BatchSpanProcessor] is being used")
             self.span_processor = BatchSpanProcessor(self.span_exporter)
         self.tag_string = tag_string
-        self.otel_service = otel_service
+        self.otel_service: str = otel_service or "airflow"
         self.resource = Resource.create(attributes={SERVICE_NAME: self.otel_service})
         self.debug = debug