Support new pycue configuration paths. (#972)

bcipriano · web-flow · commit 6e302cb9fe5e · 2021-07-29T16:21:03.000-04:00
diff --git a/pycue/opencue/config.py b/pycue/opencue/config.py
@@ -0,0 +1,96 @@
+#  Copyright Contributors to the OpenCue Project
+#
+#  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.
+
+"""OpenCue configuration."""
+
+import logging
+import os
+import platform
+
+import yaml
+
+
+logger = logging.getLogger("opencue")
+
+
+# Config file from which default settings are loaded. This file is distributed with the
+# opencue Python library.
+__DEFAULT_CONFIG_FILE = os.path.join(os.path.dirname(__file__), 'default.yaml')
+
+# Environment variables which can be used to define a custom config file. Any settings
+# defined in this file will be used instead of the defaults.
+__CONFIG_FILE_ENV_VARS = [
+    # OPENCUE_CONFIG_FILE is the preferred setting to use.
+    'OPENCUE_CONFIG_FILE',
+    # OPENCUE_CONF is deprecated, but kept for now for backwards compatibility.
+    'OPENCUE_CONF',
+]
+
+
+def config_base_directory():
+    """Returns the OpenCue config base directory.
+
+    This platform-dependent directory, stored within your user profile, is used by
+    OpenCue components as the default location for various configuration files. Typically
+    if you store your config files in this location, there is no need to set environment
+    variables to indicate where your config files are located -- OpenCue should recognize
+    them automatically.
+
+    NOTE: This work is ongoing. Over time more OpenCue components will start using this
+    base directory. See https://github.com/AcademySoftwareFoundation/OpenCue/issues/785.
+
+    :rtype: str
+    :return: config file base directory
+    """
+    if platform.system() == 'Windows':
+        return os.path.join(os.path.expandvars('%APPDATA%'), 'opencue')
+    return os.path.join(os.path.expanduser('~'), '.config', 'opencue')
+
+
+def load_config_from_file():
+    """Loads configuration settings from config file on the local system.
+
+    Default settings are read from default.yaml which is distributed with the opencue library.
+    User-provided config is then read from disk, in order of preference:
+    - Path defined by the OPENCUE_CONFIG_FILE environment variable.
+    - Path defined by the OPENCUE_CONF environment variable.
+    - Path within the config base directory (i.e. ~/.config/opencue/opencue.yaml)
+
+    :rtype: dict
+    :return: config settings
+    """
+    with open(__DEFAULT_CONFIG_FILE) as file_object:
+        config = yaml.load(file_object, Loader=yaml.SafeLoader)
+
+    user_config_file = None
+
+    for config_file_env_var in __CONFIG_FILE_ENV_VARS:
+        logger.debug('Checking for opencue config file path in %s', config_file_env_var)
+        config_file_from_env = os.environ.get(config_file_env_var)
+        if config_file_from_env and os.path.exists(config_file_from_env):
+            user_config_file = config_file_from_env
+            break
+
+    if not user_config_file:
+        config_from_user_profile = os.path.join(config_base_directory(), 'opencue.yaml')
+        logger.debug('Checking for opencue config at %s', config_from_user_profile)
+        if os.path.exists(config_from_user_profile):
+            user_config_file = config_from_user_profile
+
+    if user_config_file:
+        logger.info('Loading opencue config from %s', user_config_file)
+        with open(user_config_file) as file_object:
+            config.update(yaml.load(file_object, Loader=yaml.SafeLoader))
+
+    return config
diff --git a/pycue/opencue/cuebot.py b/pycue/opencue/cuebot.py
@@ -26,10 +26,10 @@
 import logging
 import os
 import platform
-import yaml
 
 import grpc
 
+import opencue.config
 from opencue.compiled_proto import comment_pb2
 from opencue.compiled_proto import comment_pb2_grpc
 from opencue.compiled_proto import criterion_pb2
@@ -67,15 +67,6 @@
 
 logger = logging.getLogger("opencue")
 
-default_config = os.path.join(os.path.dirname(__file__), 'default.yaml')
-with open(default_config) as file_object:
-    config = yaml.load(file_object, Loader=yaml.SafeLoader)
-
-# check for facility specific configurations.
-fcnf = os.environ.get('OPENCUE_CONF', '')
-if os.path.exists(fcnf):
-    with open(fcnf) as file_object:
-        config.update(yaml.load(file_object, Loader=yaml.SafeLoader))
 
 DEFAULT_MAX_MESSAGE_BYTES = 1024 ** 2 * 10
 DEFAULT_GRPC_PORT = 8443
@@ -96,7 +87,8 @@ class Cuebot(object):
     RpcChannel = None
     Hosts = []
     Stubs = {}
-    Timeout = config.get('cuebot.timeout', 10000)
+    Config = opencue.config.load_config_from_file()
+    Timeout = Config.get('cuebot.timeout', 10000)
 
     PROTO_MAP = {
         'action': filter_pb2,
@@ -150,13 +142,20 @@ class Cuebot(object):
     }
 
     @staticmethod
-    def init():
+    def init(config=None):
         """Main init method for setting up the Cuebot object.
-        Sets the communication channel and hosts."""
+        Sets the communication channel and hosts.
+
+        :type config: dict
+        :param config: config dictionary, this will override the config read from disk
+        """
+        if config:
+            Cuebot.Config = config
+            Cuebot.Timeout = config.get('cuebot.timeout', Cuebot.Timeout)
         if os.getenv("CUEBOT_HOSTS"):
             Cuebot.setHosts(os.getenv("CUEBOT_HOSTS").split(","))
         else:
-            facility_default = config.get("cuebot.facility_default")
+            facility_default = Cuebot.Config.get("cuebot.facility_default")
             Cuebot.setFacility(facility_default)
         if Cuebot.Hosts is None:
             raise CueException('Cuebot host not set. Please ensure CUEBOT_HOSTS is set ' +
@@ -169,7 +168,7 @@ def setChannel():
         # gRPC must specify a single host. Randomize host list to balance load across cuebots.
         hosts = list(Cuebot.Hosts)
         shuffle(hosts)
-        maxMessageBytes = config.get('cuebot.max_message_bytes', DEFAULT_MAX_MESSAGE_BYTES)
+        maxMessageBytes = Cuebot.Config.get('cuebot.max_message_bytes', DEFAULT_MAX_MESSAGE_BYTES)
 
         # create interceptors
         interceptors = (
@@ -186,7 +185,8 @@ def setChannel():
             if ':' in host:
                 connectStr = host
             else:
-                connectStr = '%s:%s' % (host, config.get('cuebot.grpc_port', DEFAULT_GRPC_PORT))
+                connectStr = '%s:%s' % (
+                    host, Cuebot.Config.get('cuebot.grpc_port', DEFAULT_GRPC_PORT))
             logger.debug('connecting to gRPC at %s', connectStr)
             # TODO(bcipriano) Configure gRPC TLS. (Issue #150)
             try:
@@ -228,12 +228,12 @@ def setFacility(facility):
 
         :type  facility: str
         :param facility: a facility named in the config file"""
-        if facility not in list(config.get("cuebot.facility").keys()):
-            default = config.get("cuebot.facility_default")
+        if facility not in list(Cuebot.Config.get("cuebot.facility").keys()):
+            default = Cuebot.Config.get("cuebot.facility_default")
             logger.warning("The facility '%s' does not exist, defaulting to %s", facility, default)
             facility = default
         logger.debug("setting facility to: %s", facility)
-        hosts = config.get("cuebot.facility")[facility]
+        hosts = Cuebot.Config.get("cuebot.facility")[facility]
         Cuebot.setHosts(hosts)
 
     @staticmethod
@@ -290,7 +290,7 @@ def getStub(cls, name):
     @staticmethod
     def getConfig():
         """Gets the Cuebot config object, originally read in from the config file on disk."""
-        return config
+        return Cuebot.Config
 
 
 # Python 2/3 compatible implementation of ABC
diff --git a/pycue/tests/config_test.py b/pycue/tests/config_test.py
@@ -0,0 +1,143 @@
+#!/usr/bin/env python
+
+#  Copyright Contributors to the OpenCue Project
+#
+#  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.
+
+"""Tests for `opencue.config`."""
+
+import os
+import unittest
+
+import mock
+import pyfakefs.fake_filesystem_unittest
+
+import opencue.config
+
+
+EXPECTED_DEFAULT_CONFIG = {
+    'logger.format': '%(levelname)-9s %(module)-10s %(message)s',
+    'logger.level': 'WARNING',
+    'cuebot.protocol': 'tcp',
+    'cuebot.grpc_port': 8443,
+    'cuebot.timeout': 10000,
+    'cuebot.max_message_bytes': 104857600,
+    'cuebot.exception_retries': 3,
+    'cuebot.facility_default': 'local',
+    'cuebot.facility': {
+        'local': ['localhost:8443'],
+        'dev': ['cuetest02-vm.example.com:8443'],
+        'cloud': [
+            'cuebot1.example.com:8443',
+            'cuebot2.example.com:8443',
+            'cuebot3.example.com:8443'
+        ],
+    },
+}
+
+USER_CONFIG = """
+cuebot.facility_default: fake-facility-01
+cuebot.facility:
+  fake-facility-01:
+    - fake-cuebot-01:1234
+  fake-facility-02:
+    - fake-cuebot-02:5678
+    - fake-cuebot-03:9012
+"""
+
+
+class ConfigTests(pyfakefs.fake_filesystem_unittest.TestCase):
+    def setUp(self):
+        self.setUpPyfakefs()
+        self.fs.add_real_file(
+            os.path.join(os.path.dirname(opencue.__file__), 'default.yaml'), read_only=True)
+        os.unsetenv('OPENCUE_CONFIG_FILE')
+        os.unsetenv('OPENCUE_CONF')
+
+    @mock.patch('platform.system', new=mock.Mock(return_value='Linux'))
+    @mock.patch('os.path.expanduser', new=mock.Mock(return_value='/home/username'))
+    def test__should_return_config_dir_unix(self):
+        self.assertEqual('/home/username/.config/opencue', opencue.config.config_base_directory())
+
+    @mock.patch('platform.system', new=mock.Mock(return_value='Windows'))
+    @mock.patch(
+        'os.path.expandvars', new=mock.Mock(return_value='C:/Users/username/AppData/Roaming'))
+    def test__should_return_config_dir_windows(self):
+        self.assertEqual(
+            'C:/Users/username/AppData/Roaming/opencue', opencue.config.config_base_directory())
+
+    def test__should_load_default_config(self):
+        self.assertIsNone(os.environ.get('OPENCUE_CONFIG_FILE'))
+        self.assertIsNone(os.environ.get('OPENCUE_CONF'))
+
+        config = opencue.config.load_config_from_file()
+
+        self.assertEqual(EXPECTED_DEFAULT_CONFIG, config)
+
+    def test__should_load_user_config(self):
+        config_file_path = '/path/to/config.yaml'
+        self.fs.create_file(config_file_path, contents=USER_CONFIG)
+        os.environ['OPENCUE_CONFIG_FILE'] = config_file_path
+        # Define some invalid config using the old setting name, this ensures the old env var
+        # will be ignored if the new one is set.
+        config_file_path_legacy = '/path/to/legacy/config.yaml'
+        self.fs.create_file(config_file_path_legacy, contents='invalid yaml')
+        os.environ['OPENCUE_CONF'] = config_file_path_legacy
+
+        config = opencue.config.load_config_from_file()
+
+        self.assertEqual('fake-facility-01', config['cuebot.facility_default'])
+        self.assertEqual(['fake-cuebot-01:1234'], config['cuebot.facility']['fake-facility-01'])
+        self.assertEqual(
+            ['fake-cuebot-02:5678', 'fake-cuebot-03:9012'],
+            config['cuebot.facility']['fake-facility-02'])
+        # Settings not defined in user config should still have default values.
+        self.assertEqual(10000, config['cuebot.timeout'])
+        self.assertEqual(3, config['cuebot.exception_retries'])
+
+    def test__should_load_user_config_from_legacy_var(self):
+        config_file_path = '/path/to/config.yaml'
+        self.fs.create_file(config_file_path, contents=USER_CONFIG)
+        os.environ['OPENCUE_CONF'] = config_file_path
+
+        config = opencue.config.load_config_from_file()
+
+        self.assertEqual('fake-facility-01', config['cuebot.facility_default'])
+        self.assertEqual(['fake-cuebot-01:1234'], config['cuebot.facility']['fake-facility-01'])
+        self.assertEqual(
+            ['fake-cuebot-02:5678', 'fake-cuebot-03:9012'],
+            config['cuebot.facility']['fake-facility-02'])
+        # Settings not defined in user config should still have default values.
+        self.assertEqual(10000, config['cuebot.timeout'])
+        self.assertEqual(3, config['cuebot.exception_retries'])
+
+    @mock.patch('platform.system', new=mock.Mock(return_value='Linux'))
+    @mock.patch('os.path.expanduser', new=mock.Mock(return_value='/home/username'))
+    def test__should_load_user_config_from_user_profile(self):
+        config_file_path = '/home/username/.config/opencue/opencue.yaml'
+        self.fs.create_file(config_file_path, contents=USER_CONFIG)
+
+        config = opencue.config.load_config_from_file()
+
+        self.assertEqual('fake-facility-01', config['cuebot.facility_default'])
+        self.assertEqual(['fake-cuebot-01:1234'], config['cuebot.facility']['fake-facility-01'])
+        self.assertEqual(
+            ['fake-cuebot-02:5678', 'fake-cuebot-03:9012'],
+            config['cuebot.facility']['fake-facility-02'])
+        # Settings not defined in user config should still have default values.
+        self.assertEqual(10000, config['cuebot.timeout'])
+        self.assertEqual(3, config['cuebot.exception_retries'])
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/pycue/tests/cuebot_test.py b/pycue/tests/cuebot_test.py
