Merge pull request #1206 from jlowin/extend-config

Add as_dict() function to print configuration

airflow/configuration.py

@@ -3,16 +3,18 @@
 from __future__ import print_function
 from __future__ import unicode_literals
 
-from future import standard_library
-standard_library.install_aliases()
-
-from builtins import str
-from configparser import ConfigParser
+import copy
 import errno
 import logging
 import os
 import subprocess
 
+from future import standard_library
+standard_library.install_aliases()
+
+from builtins import str
+from collections import OrderedDict
+from configparser import ConfigParser
 
 class AirflowConfigException(Exception):
     pass
@@ -411,29 +413,43 @@ def _validate(self):
 
         self.is_validated = True
 
+    def _get_env_var_option(self, section, key):
+        # must have format AIRFLOW__{SECTION}__{KEY} (note double underscore)
+        env_var = 'AIRFLOW__{S}__{K}'.format(S=section.upper(), K=key.upper())
+        if env_var in os.environ:
+            return expand_env_var(os.environ[env_var])
+
+    def _get_cmd_option(self, section, key):
+        fallback_key = key + '_cmd'
+        if (
+                (section, key) in ConfigParserWithDefaults.as_command_stdout and
+                self.has_option(section, fallback_key)):
+            command = self.get(section, fallback_key)
+            return run_command(command)
+
     def get(self, section, key, **kwargs):
         section = str(section).lower()
         key = str(key).lower()
-        fallback_key = key + '_cmd'
+
         d = self.defaults
 
-        # environment variables get precedence
-        # must have format AIRFLOW__{SECTION}__{KEY} (note double underscore)
-        env_var = 'AIRFLOW__{S}__{K}'.format(S=section.upper(), K=key.upper())
-        if env_var in os.environ:
-            return expand_env_var(os.environ[env_var])
+        # first check environment variables
+        option = self._get_env_var_option(section, key)
+        if option:
+            return option
 
         # ...then the config file
-        elif self.has_option(section, key):
-            return expand_env_var(ConfigParser.get(self, section, key, **kwargs))
+        if self.has_option(section, key):
+            return expand_env_var(
+                ConfigParser.get(self, section, key, **kwargs))
 
-        elif ((section, key) in ConfigParserWithDefaults.as_command_stdout
-            and self.has_option(section, fallback_key)):
-            command = self.get(section, fallback_key)
-            return run_command(command)
+        # ...then commands
+        option = self._get_cmd_option(section, key)
+        if option:
+            return option
 
         # ...then the defaults
-        elif section in d and key in d[section]:
+        if section in d and key in d[section]:
             return expand_env_var(d[section][key])
 
         else:
@@ -465,6 +481,68 @@ def read(self, filenames):
         ConfigParser.read(self, filenames)
         self._validate()
 
+    def as_dict(self, display_source=False, display_sensitive=False):
+        """
+        Returns the current configuration as an OrderedDict of OrderedDicts.
+        :param display_source: If False, the option value is returned. If True,
+            a tuple of (option_value, source) is returned. Source is either
+            'airflow.cfg' or 'default'.
+        :type display_source: bool
+        :param display_sensitive: If True, the values of options set by env
+            vars and bash commands will be displayed. If False, those options
+            are shown as '< hidden >'
+        :type display_sensitive: bool
+        """
+        cfg = copy.deepcopy(self._sections)
+
+        # remove __name__ (affects Python 2 only)
+        for options in cfg.values():
+            options.pop('__name__', None)
+
+        # add source
+        if display_source:
+            for section in cfg:
+                for k, v in cfg[section].items():
+                    cfg[section][k] = (v, 'airflow.cfg')
+
+        # add env vars and overwrite because they have priority
+        for ev in [ev for ev in os.environ if ev.startswith('AIRFLOW__')]:
+            try:
+                _, section, key = ev.split('__')
+                opt = self._get_env_var_option(section, key)
+            except ValueError:
+                opt = None
+            if opt:
+                if not display_sensitive:
+                    opt = '< hidden >'
+                if display_source:
+                    opt = (opt, 'env var')
+                cfg.setdefault(section.lower(), OrderedDict()).update(
+                    {key.lower(): opt})
+
+        # add bash commands
+        for (section, key) in ConfigParserWithDefaults.as_command_stdout:
+            opt = self._get_cmd_option(section, key)
+            if opt:
+                if not display_sensitive:
+                    opt = '< hidden >'
+                if display_source:
+                    opt = (opt, 'bash cmd')
+                cfg.setdefault(section, OrderedDict()).update({key: opt})
+
+        # add defaults
+        for section in sorted(self.defaults):
+            for key in sorted(self.defaults[section].keys()):
+                if key not in cfg.setdefault(section, OrderedDict()):
+                    opt = str(self.defaults[section][key])
+                    if display_source:
+                        cfg[section][key] = (opt, 'default')
+                    else:
+                        cfg[section][key] = opt
+
+        return cfg
+
+
 def mkdir_p(path):
     try:
         os.makedirs(path)
@@ -549,9 +627,17 @@ def getint(section, key):
 def has_option(section, key):
     return conf.has_option(section, key)
 
+
 def remove_option(section, option):
     return conf.remove_option(section, option)
 
+
+def as_dict(display_source=False, display_sensitive=False):
+    return conf.as_dict(
+        display_source=display_source, display_sensitive=display_sensitive)
+as_dict.__doc__ = conf.as_dict.__doc__
+
+
 def set(section, option, value):  # noqa
     return conf.set(section, option, value)
 

run_unit_tests.sh

@@ -1,9 +1,13 @@
 #!/bin/sh
 
+# environment
 export AIRFLOW_HOME=${AIRFLOW_HOME:=~/airflow}
 export AIRFLOW_CONFIG=$AIRFLOW_HOME/unittests.cfg
 
-# any argument received is overriding the default nose execution arguments: 
+# configuration test
+export AIRFLOW__TESTSECTION__TESTKEY=testvalue
+
+# any argument received is overriding the default nose execution arguments:
 
 nose_args=$@
 if [ -z "$nose_args" ]; then

tests/__init__.py

@@ -2,4 +2,5 @@
 from .core import *
 from .models import *
 from .operators import *
+from .configuration import *
 from .contrib import *

tests/configuration.py

@@ -0,0 +1,60 @@
+# -*- coding: utf-8 -*-
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import print_function
+import os
+import unittest
+
+from airflow import configuration
+from airflow.configuration import conf
+
+configuration.test_mode()
+
+class ConfTest(unittest.TestCase):
+    def setup(self):
+        configuration.test_mode()
+
+    def test_env_var_config(self):
+        opt = conf.get('testsection', 'testkey')
+        self.assertEqual(opt, 'testvalue')
+
+    def test_conf_as_dict(self):
+        cfg_dict = conf.as_dict()
+
+        # test that configs are picked up
+        self.assertEqual(cfg_dict['core']['unit_test_mode'], 'True')
+
+        # test env vars
+        self.assertEqual(cfg_dict['testsection']['testkey'], '< hidden >')
+
+        # test defaults
+        conf.remove_option('core', 'load_examples')
+        cfg_dict = conf.as_dict()
+        self.assertEqual(cfg_dict['core']['load_examples'], 'True')
+
+        # test display_source
+        cfg_dict = conf.as_dict(display_source=True)
+        self.assertEqual(cfg_dict['core']['unit_test_mode'][1], 'airflow.cfg')
+        self.assertEqual(cfg_dict['core']['load_examples'][1], 'default')
+        self.assertEqual(
+            cfg_dict['testsection']['testkey'], ('< hidden >', 'env var'))
+
+        # test display_sensitive
+        cfg_dict = conf.as_dict(display_sensitive=True)
+        self.assertEqual(cfg_dict['testsection']['testkey'], 'testvalue')
+
+        # test display_source and display_sensitive
+        cfg_dict = conf.as_dict(display_sensitive=True, display_source=True)
+        self.assertEqual(
+            cfg_dict['testsection']['testkey'], ('testvalue', 'env var'))