Add secure tunneling command to cli (#402)

Add new command for secure tunneling to Databricks clusters on select environments.

Author: deka108

databricks_cli/clusters/cli.py

@@ -20,7 +20,7 @@
 # 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.
-
+import sys
 import time
 from datetime import datetime
 from json import loads as json_loads
@@ -331,6 +331,41 @@ def cluster_events_cli(api_client, cluster_id, start_time, end_time, order, even
         click.echo(tabulate(_cluster_events_to_table(events_json), tablefmt='plain'))
 
 
+@click.command(context_settings=CONTEXT_SETTINGS)
+@click.option('--cluster-id', cls=OneOfOption, one_of=CLUSTER_OPTIONS,
+              type=ClusterIdClickType(), default=None, help=ClusterIdClickType.help)
+@click.option('--cluster-name', cls=OneOfOption, one_of=CLUSTER_OPTIONS,
+              type=ClusterIdClickType(), default=None, help=ClusterIdClickType.help)
+@click.option('--local-port', type=click.INT,
+              help="The local port to use for the local tunneling server")
+@click.option('--debug', '-d', is_flag=True, help="Run the tunnel in debug mode")
+@profile_option
+@eat_exceptions
+@provide_api_client
+def tunnel_cli(api_client, cluster_id, cluster_name, local_port, debug):
+    """
+    [Alpha] Start a secure TCP tunnel to a cluster over Databricks' identity proxy.
+    """
+    if sys.version_info < (3, 6):
+        raise RuntimeError("The tunneling command is not supported on Python version < 3.6")
+    if not api_client.token:
+        raise RuntimeError("The tunneling cli only supports personal token authentication.")
+
+    if cluster_id:
+        pass
+    elif cluster_name:
+        cluster = ClusterApi(api_client).get_cluster_by_name(cluster_name)
+        cluster_id = cluster["cluster_id"]
+    else:
+        raise RuntimeError('cluster_name and cluster_id must not be empty!')
+
+    # TODO(tunneling-cli): move this up once we support python3 only
+    from databricks_cli.tunnel.api import TunnelApi
+
+    click.echo("Starting a secure tunnel to cluster with ID: {}...".format(cluster_id))
+    TunnelApi(api_client, cluster_id, debug=debug).start_tunneling(local_port=local_port)
+
+
 @click.group(context_settings=CONTEXT_SETTINGS,
              short_help='Utility to interact with Databricks clusters.')
 @click.option('--version', '-v', is_flag=True, callback=print_version_callback,
@@ -358,3 +393,4 @@ def clusters_group():  # pragma: no cover
 clusters_group.add_command(spark_versions_cli, name='spark-versions')
 clusters_group.add_command(permanent_delete_cli, name='permanent-delete')
 clusters_group.add_command(cluster_events_cli, name='events')
+clusters_group.add_command(tunnel_cli, name='tunnel')

databricks_cli/commands/__init__.py

@@ -0,0 +1,51 @@
+# Databricks CLI
+# Copyright 2021 Databricks, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"), except
+# that the use of services to which certain application programming
+# interfaces (each, an "API") connect requires that the user first obtain
+# a license for the use of the APIs from Databricks, Inc. ("Databricks"),
+# by creating an account at www.databricks.com and agreeing to either (a)
+# the Community Edition Terms of Service, (b) the Databricks Terms of
+# Service, or (c) another written agreement between Licensee and Databricks
+# for the use of the APIs.
+#
+# 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 tenacity import stop_after_attempt, retry, wait_random_exponential, retry_if_result
+
+from databricks_cli.sdk.v1_service import CommandService, ExecutionContextService
+
+
+def command_is_not_terminated(resp):
+    return resp["status"] not in {"Finished", "Cancelled", "Error"}
+
+
+class CommandApi(CommandService):
+    @retry(retry=retry_if_result(command_is_not_terminated),
+           wait=wait_random_exponential(multiplier=1, max=30))
+    def wait_command_until_terminated(self, cluster_id, context_id, command_id):
+        return self.get_command_status(cluster_id, context_id, command_id)
+
+    def execute_command_until_terminated(self, language, cluster_id, context_id, command):
+        resp = self.execute_command(language=language, cluster_id=cluster_id,
+                                    context_id=context_id, command=command)
+        command_id = resp["id"]
+        return self.wait_command_until_terminated(cluster_id=cluster_id,
+                                                  context_id=context_id, command_id=command_id)
+
+
+class ExecutionContextApi(ExecutionContextService):
+    # Sometimes cluster is already in the status="RUNNING", however it couldn't provide the
+    # execution context to make the execute command stable yet.
+    @retry(stop=stop_after_attempt(10), wait=wait_random_exponential(multiplier=5, max=30))
+    def create_context(self, language, cluster_id):
+        return super().create_context(language, cluster_id)

databricks_cli/commands/api.py

@@ -104,14 +104,17 @@ def __init__(self, user=None, password=None, host=None, token=None,
         self.verify = verify
         self.api_version = api_version
         self.jobs_api_version = jobs_api_version
+        self.host = host
+        self.token = token
 
     def close(self):
         """Close the client"""
         pass
 
     # helper functions starting here
 
-    def perform_query(self, method, path, data = {}, headers = None, files=None, version=None):
+    def perform_query(self, method, path, data={}, headers=None, files=None, version=None,
+                      return_raw_response=False):
         """set up connection and perform query"""
         if headers is None:
             headers = self.default_headers
@@ -144,14 +147,16 @@ def perform_query(self, method, path, data = {}, headers = None, files=None, ver
             except ValueError:
                 pass
             raise requests.exceptions.HTTPError(message, response=e.response)
+        if return_raw_response:
+            return resp
         return resp.json()
 
 
     def get_url(self, path, version=None):
         if version:
             return self.url + version + path
         elif self.jobs_api_version and path and path.startswith('/jobs'):
-            return self.url + self.jobs_api_version + path 
+            return self.url + self.jobs_api_version + path
         return self.url + self.api_version + path
 
 

databricks_cli/sdk/api_client.py

@@ -0,0 +1,87 @@
+# Databricks CLI
+# Copyright 2021 Databricks, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"), except
+# that the use of services to which certain application programming
+# interfaces (each, an "API") connect requires that the user first obtain
+# a license for the use of the APIs from Databricks, Inc. ("Databricks"),
+# by creating an account at www.databricks.com and agreeing to either (a)
+# the Community Edition Terms of Service, (b) the Databricks Terms of
+# Service, or (c) another written agreement between Licensee and Databricks
+# for the use of the APIs.
+#
+# 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.
+import copy
+
+from databricks_cli.sdk.version import OLD_API_VERSION
+
+
+class APIV1Client(object):
+    def __init__(self, client):
+        self.v1_client = copy.deepcopy(client)
+        self.v1_client.api_version = OLD_API_VERSION
+
+
+class CommandService(APIV1Client):
+    def get_command_status(self, cluster_id, context_id, command_id):
+        return self.v1_client.perform_query(
+            method="GET", path="/commands/status", data={
+                "clusterId": cluster_id,
+                "contextId": context_id,
+                "commandId": command_id,
+            }
+        )
+
+    def cancel_command(self, cluster_id, context_id, command_id):
+        return self.v1_client.perform_query(
+            method="POST", path="/commands/cancel", data={
+                "clusterId": cluster_id,
+                "contextId": context_id,
+                "commandId": command_id,
+            }
+        )
+
+    def execute_command(self, language, cluster_id, context_id, command):
+        return self.v1_client.perform_query(
+            method="POST", path="/commands/execute", data={
+                "language": language,
+                "clusterId": cluster_id,
+                "contextId": context_id,
+                "command": command
+            }
+        )
+
+
+class ExecutionContextService(APIV1Client):
+    def get_context_status(self, cluster_id, context_id):
+        return self.v1_client.perform_query(
+            method="GET", path="/contexts/status", data={
+                "clusterId": cluster_id,
+                "contextId": context_id,
+            }
+        )
+
+    def create_context(self, language, cluster_id):
+        return self.v1_client.perform_query(
+            method="POST", path="/contexts/create", data={
+                "language": language,
+                "clusterId": cluster_id,
+            }
+        )
+
+    def destroy_context(self, cluster_id, context_id):
+        return self.v1_client.perform_query(
+            method="POST", path="/contexts/destroy", data={
+                "contextId": context_id,
+                "clusterId": cluster_id,
+            }
+        )

databricks_cli/sdk/v1_service.py

@@ -21,7 +21,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-API_VERSION='2.0'
+API_VERSION = '2.0'
+OLD_API_VERSION = '1.2'
 
 # Available API versions
-API_VERSIONS=['2.0', '2.1']
+API_VERSIONS = ['2.0', '2.1']