taskinstancekey.py 2.1 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061
  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 NamedTuple
  20. class TaskInstanceKey(NamedTuple):
  21. """Key used to identify task instance."""
  22. dag_id: str
  23. task_id: str
  24. run_id: str
  25. try_number: int = 1
  26. map_index: int = -1
  27. @property
  28. def primary(self) -> tuple[str, str, str, int]:
  29. """Return task instance primary key part of the key."""
  30. return self.dag_id, self.task_id, self.run_id, self.map_index
  31. @property
  32. def reduced(self) -> TaskInstanceKey:
  33. """Remake the key by subtracting 1 from try number to match in memory information."""
  34. # todo (dstandish): remove this property
  35. return TaskInstanceKey(
  36. self.dag_id, self.task_id, self.run_id, max(1, self.try_number - 1), self.map_index
  37. )
  38. def with_try_number(self, try_number: int) -> TaskInstanceKey:
  39. """Return TaskInstanceKey with provided ``try_number``."""
  40. return TaskInstanceKey(self.dag_id, self.task_id, self.run_id, try_number, self.map_index)
  41. @property
  42. def key(self) -> TaskInstanceKey:
  43. """
  44. For API-compatibly with TaskInstance.
  45. Returns self
  46. """
  47. return self
  48. @classmethod
  49. def from_dict(cls, dictionary):
  50. """Create TaskInstanceKey from dictionary."""
  51. return cls(**dictionary)