Create HostnameCallableRule to ease upgrade to Airflow 2.0 (apache#12743

)

Creates a rule to ensure users have a 2.0 compatible hostname_callable
configuration in their airflow.cfg

closes apache#11044

Co-authored-by: Ash Berlin-Taylor <[email protected]>
Co-authored-by: Kaxil Naik <[email protected]>

airflow/upgrade/rules/hostname_callable_rule.py

@@ -0,0 +1,38 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you 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 absolute_import
+
+from airflow.configuration import conf
+from airflow.upgrade.rules.base_rule import BaseRule
+
+
+class HostnameCallable(BaseRule):
+    title = "Unify hostname_callable option in core section"
+
+    description = ""
+
+    def check(self):
+        hostname_callable_conf = conf.get("core", "hostname_callable")
+        if ":" in hostname_callable_conf:
+            return (
+                "Error: hostname_callable `{}` "
+                "contains a colon instead of a dot. please change to `{}`".format(
+                    hostname_callable_conf, hostname_callable_conf.replace(":", ".")
+                )
+            )
+        return None

airflow/utils/net.py

@@ -43,7 +43,10 @@ def get_hostname():
         return socket.getfqdn()
 
     # Since we have a callable path, we try to import and run it next.
-    module_path, attr_name = callable_path.split(':')
-    module = importlib.import_module(module_path)
-    callable = getattr(module, attr_name)
-    return callable()
+    if ":" in callable_path:
+        module_path, attr_name = callable_path.split(':')
+        module = importlib.import_module(module_path)
+        callable = getattr(module, attr_name)
+        return callable()
+    else:
+        return conf.getimport('core', 'hostname_callable', fallback='socket.getfqdn')()

tests/upgrade/rules/test_hostname_callable_rule.py

@@ -0,0 +1,36 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you 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 unittest import TestCase
+
+from airflow.upgrade.rules.hostname_callable_rule import HostnameCallable
+from tests.test_utils.config import conf_vars
+
+
+class TestFernetEnabledRule(TestCase):
+    @conf_vars({("core", "hostname_callable"): "dummyhostname:function"})
+    def test_incorrect_hostname(self):
+        result = HostnameCallable().check()
+        self.assertEqual(
+            result,
+            "Error: hostname_callable `dummyhostname:function` "
+            "contains a colon instead of a dot. please change to `dummyhostname.function`",
+        )
+
+    @conf_vars({("core", "hostname_callable"): "dummyhostname.function"})
+    def test_correct_hostname(self):
+        result = HostnameCallable().check()
+        self.assertEqual(result, None)

tests/utils/test_net.py

@@ -19,8 +19,12 @@
 
 import unittest
 import mock
+import six
 
 from airflow.utils import net
+from tests.test_utils.config import conf_vars
+from airflow.exceptions import AirflowConfigException
+import re
 
 
 def get_hostname():
@@ -43,12 +47,9 @@ def test_get_hostname_set(self, patched_conf):
         )
         self.assertTrue(net.get_hostname() == 'awesomehostname')
 
-    @mock.patch('airflow.utils.net.conf')
-    def test_get_hostname_set_incorrect(self, patched_conf):
-        patched_conf.get = mock.Mock(
-            return_value='tests.utils.test_net'
-        )
-        with self.assertRaises(ValueError):
+    @conf_vars({('core', 'hostname_callable'): 'tests.utils.test_net'})
+    def test_get_hostname_set_incorrect(self):
+        with self.assertRaises(TypeError):
             net.get_hostname()
 
     @mock.patch('airflow.utils.net.conf')
@@ -58,3 +59,29 @@ def test_get_hostname_set_missing(self, patched_conf):
         )
         with self.assertRaises(AttributeError):
             net.get_hostname()
+
+    @mock.patch('socket.getfqdn', return_value='first')
+    @conf_vars({('core', 'hostname_callable'): None})
+    def test_get_hostname_unset_2_0(self, mock_getfqdn):
+        self.assertEqual('first', net.get_hostname())
+
+    @conf_vars({('core', 'hostname_callable'): 'tests.utils.test_net.get_hostname'})
+    def test_get_hostname_set_2_0(self):
+        self.assertEqual('awesomehostname', net.get_hostname())
+
+    @conf_vars({('core', 'hostname_callable'): 'tests.utils.test_net'})
+    def test_get_hostname_set_incorrect_2_0(self):
+        with self.assertRaises(TypeError):
+            net.get_hostname()
+
+    @conf_vars({('core', 'hostname_callable'): 'tests.utils.test_net.missing_func'})
+    def test_get_hostname_set_missing_2_0(self):
+        with six.assertRaisesRegex(
+            self,
+            AirflowConfigException,
+            re.escape(
+                'The object could not be loaded. Please check "hostname_callable" key in "core" section. '
+                'Current value: "tests.utils.test_net.missing_func"'
+            ),
+        ):
+            net.get_hostname()