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

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, relationship 

24 

25from airflow.models.base import Base, StringID 

26from airflow.utils.session import NEW_SESSION, provide_session 

27from airflow.utils.sqlalchemy import mapped_column 

28 

29if TYPE_CHECKING: 

30 from sqlalchemy.orm import Session 

31 

32dag_bundle_team_association_table = Table( 

33 "dag_bundle_team", 

34 Base.metadata, 

35 Column( 

36 "dag_bundle_name", 

37 StringID(length=250), 

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

39 primary_key=True, 

40 ), 

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

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

43 Index("idx_dag_bundle_team_team_name", "team_name"), 

44) 

45 

46 

47class Team(Base): 

48 """ 

49 Contains the list of teams defined in the environment. 

50 

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

52 """ 

53 

54 __tablename__ = "team" 

55 

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

57 dag_bundles = relationship( 

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

59 ) 

60 

61 def __repr__(self): 

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

63 

64 @classmethod 

65 @provide_session 

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

67 """ 

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

69 

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

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

72 

73 :return: Set of all team names 

74 """ 

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