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(self, conn_id: str, session: Session = NEW_SESSION) -> Connection | None: 

40 """ 

41 Get Airflow Connection from Metadata DB. 

42 

43 :param conn_id: Connection ID 

44 :param session: SQLAlchemy Session 

45 :return: Connection Object 

46 """ 

47 from airflow.models import Connection 

48 

49 conn = session.scalar(select(Connection).where(Connection.conn_id == conn_id).limit(1)) 

50 session.expunge_all() 

51 return conn 

52 

53 @provide_session 

54 def get_variable( 

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

56 ) -> str | None: 

57 """ 

58 Get Airflow Variable from Metadata DB. 

59 

60 :param key: Variable Key 

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

62 :param session: SQLAlchemy Session 

63 :return: Variable Value 

64 """ 

65 from airflow.models import Variable 

66 

67 var_value = session.scalar( 

68 select(Variable) 

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

70 .limit(1) 

71 ) 

72 session.expunge_all() 

73 if var_value: 

74 return var_value.val 

75 return None