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

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.""" 

19from __future__ import annotations 

20 

21import warnings 

22from typing import TYPE_CHECKING 

23 

24from sqlalchemy.orm import Session 

25 

26from airflow.exceptions import RemovedInAirflow3Warning 

27from airflow.secrets import BaseSecretsBackend 

28from airflow.utils.session import NEW_SESSION, provide_session 

29 

30if TYPE_CHECKING: 

31 from airflow.models.connection import Connection 

32 

33 

34class MetastoreBackend(BaseSecretsBackend): 

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

36 

37 @provide_session 

38 def get_connection(self, conn_id: str, session: Session = NEW_SESSION) -> Connection | None: 

39 from airflow.models.connection import Connection 

40 

41 conn = session.query(Connection).filter(Connection.conn_id == conn_id).first() 

42 session.expunge_all() 

43 return conn 

44 

45 @provide_session 

46 def get_connections(self, conn_id: str, session: Session = NEW_SESSION) -> list[Connection]: 

47 warnings.warn( 

48 "This method is deprecated. Please use " 

49 "`airflow.secrets.metastore.MetastoreBackend.get_connection`.", 

50 RemovedInAirflow3Warning, 

51 stacklevel=3, 

52 ) 

53 conn = self.get_connection(conn_id=conn_id, session=session) 

54 if conn: 

55 return [conn] 

56 return [] 

57 

58 @provide_session 

59 def get_variable(self, key: str, session: Session = NEW_SESSION) -> str | None: 

60 """ 

61 Get Airflow Variable from Metadata DB. 

62 

63 :param key: Variable Key 

64 :return: Variable Value 

65 """ 

66 from airflow.models.variable import Variable 

67 

68 var_value = session.query(Variable).filter(Variable.key == key).first() 

69 session.expunge_all() 

70 if var_value: 

71 return var_value.val 

72 return None