Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/airflow/secrets/environment_variables.py: 64%

Shortcuts on this page

r m x   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

14 statements  

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

19 

20from __future__ import annotations 

21 

22import os 

23 

24from airflow.secrets import BaseSecretsBackend 

25 

26CONN_ENV_PREFIX = "AIRFLOW_CONN_" 

27VAR_ENV_PREFIX = "AIRFLOW_VAR_" 

28 

29 

30class EnvironmentVariablesBackend(BaseSecretsBackend): 

31 """Retrieves Connection object and Variable from environment variable.""" 

32 

33 def get_conn_value(self, conn_id: str) -> str | None: 

34 return os.environ.get(CONN_ENV_PREFIX + conn_id.upper()) 

35 

36 def get_variable(self, key: str, team_name: str | None = None) -> str | None: 

37 """ 

38 Get Airflow Variable from Environment Variable. 

39 

40 :param key: Variable Key 

41 :param team_name: Team name associated to the task trying to access the variable (if any) 

42 :return: Variable Value 

43 """ 

44 if team_name and ( 

45 team_var := os.environ.get(f"{VAR_ENV_PREFIX}_{team_name.upper()}___" + key.upper()) 

46 ): 

47 # Format to set a team specific variable: AIRFLOW_VAR__<TEAM_ID>___<VAR_KEY> 

48 return team_var 

49 

50 return os.environ.get(VAR_ENV_PREFIX + key.upper())