Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/airflow/secrets/metastore.py: 48%

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"""Objects relating to sourcing connections from metastore database.""" 

19 

20from __future__ import annotations 

21 

22from typing import TYPE_CHECKING 

23 

24from sqlalchemy import or_, select 

25 

26from airflow.secrets import BaseSecretsBackend 

27from airflow.utils.session import NEW_SESSION, provide_session 

28 

29if TYPE_CHECKING: 

30 from sqlalchemy.orm import Session 

31 

32 from airflow.models import Connection 

33 

34 

35class MetastoreBackend(BaseSecretsBackend): 

36 """Retrieves Connection object and Variable from airflow metastore database.""" 

37 

38 @provide_session 

39 def get_connection( 

40 self, conn_id: str, team_name: str | None = None, session: Session = NEW_SESSION 

41 ) -> Connection | None: 

42 """ 

43 Get Airflow Connection from Metadata DB. 

44 

45 :param conn_id: Connection ID 

46 :param team_name: Team name associated to the task trying to access the connection (if any) 

47 :param session: SQLAlchemy Session 

48 :return: Connection Object 

49 """ 

50 from airflow.models import Connection 

51 

52 conn = session.scalar( 

53 select(Connection) 

54 .where( 

55 Connection.conn_id == conn_id, 

56 or_(Connection.team_name == team_name, Connection.team_name.is_(None)), 

57 ) 

58 .limit(1) 

59 ) 

60 session.expunge_all() 

61 return conn 

62 

63 @provide_session 

64 def get_variable( 

65 self, key: str, team_name: str | None = None, session: Session = NEW_SESSION 

66 ) -> str | None: 

67 """ 

68 Get Airflow Variable from Metadata DB. 

69 

70 :param key: Variable Key 

71 :param team_name: Team name associated to the task trying to access the variable (if any) 

72 :param session: SQLAlchemy Session 

73 :return: Variable Value 

74 """ 

75 from airflow.models import Variable 

76 

77 var_value = session.scalar( 

78 select(Variable) 

79 .where(Variable.key == key, or_(Variable.team_name == team_name, Variable.team_name.is_(None))) 

80 .limit(1) 

81 ) 

82 session.expunge_all() 

83 if var_value: 

84 return var_value.val 

85 return None