Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/airflow/models/team.py: 80%

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 TYPE_CHECKING 

21 

22from sqlalchemy import Column, ForeignKey, Index, String, Table, select 

23from sqlalchemy.orm import Mapped, mapped_column, relationship 

24 

25from airflow.models.base import Base, StringID 

26from airflow.utils.session import NEW_SESSION, provide_session 

27 

28if TYPE_CHECKING: 

29 from sqlalchemy.orm import Session 

30 

31dag_bundle_team_association_table = Table( 

32 "dag_bundle_team", 

33 Base.metadata, 

34 Column( 

35 "dag_bundle_name", 

36 StringID(length=250), 

37 ForeignKey("dag_bundle.name", ondelete="CASCADE"), 

38 primary_key=True, 

39 ), 

40 Column("team_name", String(50), ForeignKey("team.name", ondelete="CASCADE"), primary_key=True), 

41 Index("idx_dag_bundle_team_dag_bundle_name", "dag_bundle_name", unique=True), 

42 Index("idx_dag_bundle_team_team_name", "team_name"), 

43) 

44 

45 

46class Team(Base): 

47 """ 

48 Contains the list of teams defined in the environment. 

49 

50 This table is only used when Airflow is run in multi-team mode. 

51 """ 

52 

53 __tablename__ = "team" 

54 

55 name: Mapped[str] = mapped_column(String(50), primary_key=True) 

56 dag_bundles = relationship( 

57 "DagBundleModel", secondary=dag_bundle_team_association_table, back_populates="teams" 

58 ) 

59 

60 def __repr__(self): 

61 return f"Team(name={self.name})" 

62 

63 @classmethod 

64 @provide_session 

65 def get_all_team_names(cls, session: Session = NEW_SESSION) -> set[str]: 

66 """ 

67 Return a set of all team names from the database. 

68 

69 This method provides a convenient way to get just the team names for validation 

70 purposes, such as verifying team names in executor configurations. 

71 

72 :return: Set of all team names 

73 """ 

74 return set(session.scalars(select(Team.name)).all())