Coverage for /pythoncovmergedfiles/medio/medio/src/airflow/build/lib/airflow/models/taskinstancekey.py: 79%

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. 

18from __future__ import annotations 

19 

20from typing import NamedTuple 

21 

22 

23class TaskInstanceKey(NamedTuple): 

24 """Key used to identify task instance.""" 

25 

26 dag_id: str 

27 task_id: str 

28 run_id: str 

29 try_number: int = 1 

30 map_index: int = -1 

31 

32 @property 

33 def primary(self) -> tuple[str, str, str, int]: 

34 """Return task instance primary key part of the key.""" 

35 return self.dag_id, self.task_id, self.run_id, self.map_index 

36 

37 @property 

38 def reduced(self) -> TaskInstanceKey: 

39 """Remake the key by subtracting 1 from try number to match in memory information.""" 

40 return TaskInstanceKey( 

41 self.dag_id, self.task_id, self.run_id, max(1, self.try_number - 1), self.map_index 

42 ) 

43 

44 def with_try_number(self, try_number: int) -> TaskInstanceKey: 

45 """Returns TaskInstanceKey with provided ``try_number``.""" 

46 return TaskInstanceKey(self.dag_id, self.task_id, self.run_id, try_number, self.map_index) 

47 

48 @property 

49 def key(self) -> TaskInstanceKey: 

50 """For API-compatibly with TaskInstance. 

51 

52 Returns self 

53 """ 

54 return self