base_stats_logger.py 3.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114
  1. # Licensed to the Apache Software Foundation (ASF) under one
  2. # or more contributor license agreements. See the NOTICE file
  3. # distributed with this work for additional information
  4. # regarding copyright ownership. The ASF licenses this file
  5. # to you under the Apache License, Version 2.0 (the
  6. # "License"); you may not use this file except in compliance
  7. # with the License. You may obtain a copy of the License at
  8. #
  9. # http://www.apache.org/licenses/LICENSE-2.0
  10. #
  11. # Unless required by applicable law or agreed to in writing,
  12. # software distributed under the License is distributed on an
  13. # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  14. # KIND, either express or implied. See the License for the
  15. # specific language governing permissions and limitations
  16. # under the License.
  17. from __future__ import annotations
  18. from typing import TYPE_CHECKING, Any
  19. from airflow.metrics.protocols import Timer
  20. from airflow.typing_compat import Protocol
  21. if TYPE_CHECKING:
  22. from airflow.metrics.protocols import DeltaType, TimerProtocol
  23. class StatsLogger(Protocol):
  24. """This class is only used for TypeChecking (for IDEs, mypy, etc)."""
  25. instance: StatsLogger | NoStatsLogger | None = None
  26. @classmethod
  27. def incr(
  28. cls,
  29. stat: str,
  30. count: int = 1,
  31. rate: int | float = 1,
  32. *,
  33. tags: dict[str, Any] | None = None,
  34. ) -> None:
  35. """Increment stat."""
  36. @classmethod
  37. def decr(
  38. cls,
  39. stat: str,
  40. count: int = 1,
  41. rate: int | float = 1,
  42. *,
  43. tags: dict[str, Any] | None = None,
  44. ) -> None:
  45. """Decrement stat."""
  46. @classmethod
  47. def gauge(
  48. cls,
  49. stat: str,
  50. value: float,
  51. rate: int | float = 1,
  52. delta: bool = False,
  53. *,
  54. tags: dict[str, Any] | None = None,
  55. ) -> None:
  56. """Gauge stat."""
  57. @classmethod
  58. def timing(
  59. cls,
  60. stat: str,
  61. dt: DeltaType | None,
  62. *,
  63. tags: dict[str, Any] | None = None,
  64. ) -> None:
  65. """Stats timing."""
  66. @classmethod
  67. def timer(cls, *args, **kwargs) -> TimerProtocol:
  68. """Timer metric that can be cancelled."""
  69. raise NotImplementedError()
  70. class NoStatsLogger:
  71. """If no StatsLogger is configured, NoStatsLogger is used as a fallback."""
  72. @classmethod
  73. def incr(cls, stat: str, count: int = 1, rate: int = 1, *, tags: dict[str, str] | None = None) -> None:
  74. """Increment stat."""
  75. @classmethod
  76. def decr(cls, stat: str, count: int = 1, rate: int = 1, *, tags: dict[str, str] | None = None) -> None:
  77. """Decrement stat."""
  78. @classmethod
  79. def gauge(
  80. cls,
  81. stat: str,
  82. value: int,
  83. rate: int = 1,
  84. delta: bool = False,
  85. *,
  86. tags: dict[str, str] | None = None,
  87. ) -> None:
  88. """Gauge stat."""
  89. @classmethod
  90. def timing(cls, stat: str, dt: DeltaType, *, tags: dict[str, str] | None = None) -> None:
  91. """Stats timing."""
  92. @classmethod
  93. def timer(cls, *args, **kwargs) -> TimerProtocol:
  94. """Timer metric that can be cancelled."""
  95. return Timer()