Add a CredentialProvider reading from TaskContext in Spark Executors. (#232)

* Added a credential provider for accessing crdentials form the TaskContext on Spark Executors.

* Added pylint directive to ignore broken import.

* Adressed review comments.

* Adressed review comments - added test making sure new provider does not break others if not available.

* Adressed review comment.

* Fixed linter complaints.

* Addressed linter complaints.

* Addressed linter complaints.

* Addressed linter complaints.

* Addressed linter complaints.

* Added passing of ignoreTls flag to the executers.

* Fixed typo.

* nit

* lint

* lint

* Renamed method.

* set insecure classmethod -> staticmethod

Author: tomasatdatabricks

databricks_cli/configure/provider.py

@@ -185,19 +185,49 @@ def get_config(self):
 
 
 class DefaultConfigProvider(DatabricksConfigProvider):
-    """Prefers environment variables, and then the default profile."""
+    """Look for credentials in a chain of default locations."""
     def __init__(self):
-        self.env_provider = EnvironmentVariableConfigProvider()
-        self.default_profile_provider = ProfileConfigProvider()
+        self._providers = (
+            SparkTaskContextConfigProvider(),
+            EnvironmentVariableConfigProvider(),
+            ProfileConfigProvider()
+        )
 
     def get_config(self):
-        env_config = self.env_provider.get_config()
-        if env_config:
-            return env_config
+        for provider in self._providers:
+            config = provider.get_config()
+            if config is not None and config.is_valid:
+                return config
+        return None
+
+
+class SparkTaskContextConfigProvider(DatabricksConfigProvider):
+    """Loads credentials from Spark TaskContext if running in a Spark Executor."""
+
+    @staticmethod
+    def _get_spark_task_context_or_none():
+        try:
+            from pyspark import TaskContext  # pylint: disable=import-error
+            return TaskContext.get()
+        except ImportError:
+            return None
 
-        profile_config = self.default_profile_provider.get_config()
-        if profile_config:
-            return profile_config
+    @staticmethod
+    def set_insecure(x):
+        from pyspark import SparkContext  # pylint: disable=import-error      
+        new_val = "True" if x else None
+        SparkContext._active_spark_context.setLocalProperty("spark.databricks.ignoreTls", new_val)
+
+    def get_config(self):
+        context = self._get_spark_task_context_or_none()
+        if context is not None:
+            host = context.getLocalProperty("spark.databricks.api.url")
+            token = context.getLocalProperty("spark.databricks.token")
+            insecure = context.getLocalProperty("spark.databricks.ignoreTls")
+            config = DatabricksConfig.from_token(host=host, token=token, insecure=insecure)
+            if config.is_valid:
+                return config
+        return None
 
 
 class EnvironmentVariableConfigProvider(DatabricksConfigProvider):

tests/configure/test_provider.py

@@ -28,7 +28,8 @@
 
 from databricks_cli.configure.provider import DatabricksConfig, DEFAULT_SECTION, \
     update_and_persist_config, get_config_for_profile, get_config, \
-    set_config_provider, ProfileConfigProvider, _get_path, DatabricksConfigProvider
+    set_config_provider, ProfileConfigProvider, _get_path, DatabricksConfigProvider,\
+    SparkTaskContextConfigProvider
 from databricks_cli.utils import InvalidConfigurationError
 
 
@@ -133,6 +134,38 @@ def test_get_config_uses_default_profile():
     assert config.token == "hello"
 
 
+def test_get_config_uses_task_context_variable():
+    class TaskContextMock(object):
+
+        def __init__(self):
+            pass
+
+        def getLocalProperty(self, x):  # NOQA
+            if x == "spark.databricks.api.url":
+                return "url"
+            elif x == "spark.databricks.token":
+                return "token"
+            elif x == "spark.databricks.ignoreTls":
+                return "True"
+            else:
+                raise Exception("should not get here.")
+
+    ctx_class = ("databricks_cli.configure.provider.SparkTaskContextConfigProvider."
+                 "_get_spark_task_context_or_none")
+    with patch(ctx_class) as get_context_mock:
+        get_context_mock.return_value = TaskContextMock()
+        config = get_config()
+        assert config.host == "url"
+        assert config.token == "token"
+        assert config.insecure == "True"
+        assert config.username is None
+        assert config.password is None
+
+
+def test_task_context_provider_does_not_break_stuff():
+    assert SparkTaskContextConfigProvider().get_config() is None
+
+
 def test_get_config_uses_env_variable():
     with patch.dict('os.environ', {'DATABRICKS_HOST': TEST_HOST,
                                    'DATABRICKS_USERNAME': TEST_USER,