log.py 3.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102
  1. #
  2. # Licensed to the Apache Software Foundation (ASF) under one
  3. # or more contributor license agreements. See the NOTICE file
  4. # distributed with this work for additional information
  5. # regarding copyright ownership. The ASF licenses this file
  6. # to you under the Apache License, Version 2.0 (the
  7. # "License"); you may not use this file except in compliance
  8. # with the License. You may obtain a copy of the License at
  9. #
  10. # http://www.apache.org/licenses/LICENSE-2.0
  11. #
  12. # Unless required by applicable law or agreed to in writing,
  13. # software distributed under the License is distributed on an
  14. # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  15. # KIND, either express or implied. See the License for the
  16. # specific language governing permissions and limitations
  17. # under the License.
  18. from __future__ import annotations
  19. from typing import TYPE_CHECKING
  20. from sqlalchemy import Column, Index, Integer, String, Text
  21. from airflow.models.base import Base, StringID
  22. from airflow.utils import timezone
  23. from airflow.utils.sqlalchemy import UtcDateTime
  24. if TYPE_CHECKING:
  25. from airflow.models.taskinstance import TaskInstance
  26. from airflow.models.taskinstancekey import TaskInstanceKey
  27. class Log(Base):
  28. """Used to actively log events to the database."""
  29. __tablename__ = "log"
  30. id = Column(Integer, primary_key=True)
  31. dttm = Column(UtcDateTime)
  32. dag_id = Column(StringID())
  33. task_id = Column(StringID())
  34. map_index = Column(Integer)
  35. event = Column(String(60))
  36. execution_date = Column(UtcDateTime)
  37. run_id = Column(StringID())
  38. owner = Column(String(500))
  39. owner_display_name = Column(String(500))
  40. extra = Column(Text)
  41. try_number = Column(Integer)
  42. __table_args__ = (
  43. Index("idx_log_dag", dag_id),
  44. Index("idx_log_dttm", dttm),
  45. Index("idx_log_event", event),
  46. Index("idx_log_task_instance", dag_id, task_id, run_id, map_index, try_number),
  47. )
  48. def __init__(
  49. self,
  50. event,
  51. task_instance: TaskInstance | TaskInstanceKey | None = None,
  52. owner=None,
  53. owner_display_name=None,
  54. extra=None,
  55. **kwargs,
  56. ):
  57. self.dttm = timezone.utcnow()
  58. self.event = event
  59. self.extra = extra
  60. task_owner = None
  61. self.execution_date = None
  62. if task_instance:
  63. self.dag_id = task_instance.dag_id
  64. self.task_id = task_instance.task_id
  65. if execution_date := getattr(task_instance, "execution_date", None):
  66. self.execution_date = execution_date
  67. self.run_id = task_instance.run_id
  68. self.try_number = task_instance.try_number
  69. self.map_index = task_instance.map_index
  70. if task := getattr(task_instance, "task", None):
  71. task_owner = task.owner
  72. if "task_id" in kwargs:
  73. self.task_id = kwargs["task_id"]
  74. if "dag_id" in kwargs:
  75. self.dag_id = kwargs["dag_id"]
  76. if kwargs.get("execution_date"):
  77. self.execution_date = kwargs["execution_date"]
  78. if kwargs.get("run_id"):
  79. self.run_id = kwargs["run_id"]
  80. if "map_index" in kwargs:
  81. self.map_index = kwargs["map_index"]
  82. if "try_number" in kwargs:
  83. self.try_number = kwargs["try_number"]
  84. self.owner = owner or task_owner
  85. self.owner_display_name = owner_display_name or None
  86. def __str__(self) -> str:
  87. return f"Log({self.event}, {self.task_id}, {self.owner}, {self.owner_display_name}, {self.extra})"