Skip to content

Commit 2158dc4

Browse files
committed
remove internal_api_call from airflow.secrets.metastore
Part of #44436
1 parent 13e5464 commit 2158dc4

File tree

1 file changed

+14
-20
lines changed

1 file changed

+14
-20
lines changed

airflow/secrets/metastore.py

Lines changed: 14 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -23,49 +23,43 @@
2323

2424
from sqlalchemy import select
2525

26-
from airflow.api_internal.internal_api_call import internal_api_call
26+
from airflow.models import Variable
2727
from airflow.secrets import BaseSecretsBackend
2828
from airflow.utils.session import NEW_SESSION, provide_session
2929

3030
if TYPE_CHECKING:
3131
from sqlalchemy.orm import Session
3232

33-
from airflow.models.connection import Connection
33+
from airflow.models import Connection
3434

3535

3636
class MetastoreBackend(BaseSecretsBackend):
3737
"""Retrieves Connection object and Variable from airflow metastore database."""
3838

3939
@provide_session
4040
def get_connection(self, conn_id: str, session: Session = NEW_SESSION) -> Connection | None:
41-
return MetastoreBackend._fetch_connection(conn_id, session=session)
42-
43-
@provide_session
44-
def get_variable(self, key: str, session: Session = NEW_SESSION) -> str | None:
4541
"""
46-
Get Airflow Variable from Metadata DB.
42+
Get Airflow Connection from Metadata DB.
4743
48-
:param key: Variable Key
49-
:return: Variable Value
44+
:param conn_id: Connection ID
45+
:param session: SQLAlchemy Session
46+
:return: Connection Object
5047
"""
51-
return MetastoreBackend._fetch_variable(key=key, session=session)
52-
53-
@staticmethod
54-
@internal_api_call
55-
@provide_session
56-
def _fetch_connection(conn_id: str, session: Session = NEW_SESSION) -> Connection | None:
57-
from airflow.models.connection import Connection
48+
from airflow.models import Connection
5849

5950
conn = session.scalar(select(Connection).where(Connection.conn_id == conn_id).limit(1))
6051
session.expunge_all()
6152
return conn
6253

63-
@staticmethod
64-
@internal_api_call
6554
@provide_session
66-
def _fetch_variable(key: str, session: Session = NEW_SESSION) -> str | None:
67-
from airflow.models.variable import Variable
55+
def get_variable(self, key: str, session: Session = NEW_SESSION) -> str | None:
56+
"""
57+
Get Airflow Variable from Metadata DB.
6858
59+
:param key: Variable Key
60+
:param session: SQLAlchemy Session
61+
:return: Variable Value
62+
"""
6963
var_value = session.scalar(select(Variable).where(Variable.key == key).limit(1))
7064
session.expunge_all()
7165
if var_value:

0 commit comments

Comments
 (0)