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

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 

22import warnings 

23from typing import TYPE_CHECKING 

24 

25from sqlalchemy import select 

26 

27from airflow.api_internal.internal_api_call import internal_api_call 

28from airflow.exceptions import RemovedInAirflow3Warning 

29from airflow.secrets import BaseSecretsBackend 

30from airflow.utils.session import NEW_SESSION, provide_session 

31 

32if TYPE_CHECKING: 

33 from sqlalchemy.orm import Session 

34 

35 from airflow.models.connection import Connection 

36 

37 

38class MetastoreBackend(BaseSecretsBackend): 

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

40 

41 @provide_session 

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

43 return MetastoreBackend._fetch_connection(conn_id, session=session) 

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 return MetastoreBackend._fetch_variable(key=key, session=session) 

67 

68 @staticmethod 

69 @internal_api_call 

70 @provide_session 

71 def _fetch_connection(conn_id: str, session: Session = NEW_SESSION) -> Connection | None: 

72 from airflow.models.connection import Connection 

73 

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

75 session.expunge_all() 

76 return conn 

77 

78 @staticmethod 

79 @internal_api_call 

80 @provide_session 

81 def _fetch_variable(key: str, session: Session = NEW_SESSION) -> str | None: 

82 from airflow.models.variable import Variable 

83 

84 var_value = session.scalar(select(Variable).where(Variable.key == key).limit(1)) 

85 session.expunge_all() 

86 if var_value: 

87 return var_value.val 

88 return None