Coverage for /pythoncovmergedfiles/medio/medio/src/airflow/build/lib/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.""" 

19from __future__ import annotations 

20 

21import warnings 

22from typing import TYPE_CHECKING 

23 

24from airflow.exceptions import RemovedInAirflow3Warning 

25from airflow.secrets import BaseSecretsBackend 

26from airflow.utils.session import provide_session 

27 

28if TYPE_CHECKING: 

29 from airflow.models.connection import Connection 

30 

31 

32class MetastoreBackend(BaseSecretsBackend): 

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

34 

35 @provide_session 

36 def get_connection(self, conn_id, session=None) -> Connection | None: 

37 from airflow.models.connection import Connection 

38 

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

40 session.expunge_all() 

41 return conn 

42 

43 @provide_session 

44 def get_connections(self, conn_id, session=None) -> list[Connection]: 

45 warnings.warn( 

46 "This method is deprecated. Please use " 

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

48 RemovedInAirflow3Warning, 

49 stacklevel=3, 

50 ) 

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

52 if conn: 

53 return [conn] 

54 return [] 

55 

56 @provide_session 

57 def get_variable(self, key: str, session=None): 

58 """ 

59 Get Airflow Variable from Metadata DB. 

60 

61 :param key: Variable Key 

62 :return: Variable Value 

63 """ 

64 from airflow.models.variable import Variable 

65 

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

67 session.expunge_all() 

68 if var_value: 

69 return var_value.val 

70 return None