# 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. from __future__ import annotations from typing import TYPE_CHECKING, Any from airflow.metrics.protocols import Timer from airflow.typing_compat import Protocol if TYPE_CHECKING: from airflow.metrics.protocols import DeltaType, TimerProtocol class StatsLogger(Protocol): """This class is only used for TypeChecking (for IDEs, mypy, etc).""" instance: StatsLogger | NoStatsLogger | None = None @classmethod def incr( cls, stat: str, count: int = 1, rate: int | float = 1, *, tags: dict[str, Any] | None = None, ) -> None: """Increment stat.""" @classmethod def decr( cls, stat: str, count: int = 1, rate: int | float = 1, *, tags: dict[str, Any] | None = None, ) -> None: """Decrement stat.""" @classmethod def gauge( cls, stat: str, value: float, rate: int | float = 1, delta: bool = False, *, tags: dict[str, Any] | None = None, ) -> None: """Gauge stat.""" @classmethod def timing( cls, stat: str, dt: DeltaType | None, *, tags: dict[str, Any] | None = None, ) -> None: """Stats timing.""" @classmethod def timer(cls, *args, **kwargs) -> TimerProtocol: """Timer metric that can be cancelled.""" raise NotImplementedError() class NoStatsLogger: """If no StatsLogger is configured, NoStatsLogger is used as a fallback.""" @classmethod def incr(cls, stat: str, count: int = 1, rate: int = 1, *, tags: dict[str, str] | None = None) -> None: """Increment stat.""" @classmethod def decr(cls, stat: str, count: int = 1, rate: int = 1, *, tags: dict[str, str] | None = None) -> None: """Decrement stat.""" @classmethod def gauge( cls, stat: str, value: int, rate: int = 1, delta: bool = False, *, tags: dict[str, str] | None = None, ) -> None: """Gauge stat.""" @classmethod def timing(cls, stat: str, dt: DeltaType, *, tags: dict[str, str] | None = None) -> None: """Stats timing.""" @classmethod def timer(cls, *args, **kwargs) -> TimerProtocol: """Timer metric that can be cancelled.""" return Timer()