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