Merge pull request #93 from keboola/snowflake-keypair-workspace-default

feat: default to key-pair login for Snowflake workspaces

Author: jirkasemmler

kbcstorage/workspaces.py

@@ -6,12 +6,45 @@
 .. _here:
     http://docs.keboola.apiary.io/#reference/workspaces/
 """
+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric import rsa
+
 from kbcstorage.base import Endpoint
 from kbcstorage.files import Files
 from kbcstorage.jobs import Jobs
+from kbcstorage.tokens import Tokens
 from typing import List  # the legacy Workspaces class below unfortunately defines its own method called list
 
 
+BACKEND_SNOWFLAKE = 'snowflake'
+LOGIN_TYPE_DEFAULT = 'default'
+LOGIN_TYPE_SNOWFLAKE_SERVICE_KEYPAIR = 'snowflake-service-keypair'
+
+# sentinel distinguishing "not resolved yet" from "project has no default backend"
+_DEFAULT_BACKEND_UNRESOLVED = object()
+
+
+def _generate_rsa_key_pair():
+    """
+    Generate an RSA-2048 key pair for Snowflake key-pair authentication.
+
+    Returns:
+        (private_key_pem, public_key_pem): Both keys PEM-encoded, the private
+            key in PKCS#8 format as expected by Snowflake drivers.
+    """
+    private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
+    private_key_pem = private_key.private_bytes(
+        encoding=serialization.Encoding.PEM,
+        format=serialization.PrivateFormat.PKCS8,
+        encryption_algorithm=serialization.NoEncryption(),
+    ).decode('ascii')
+    public_key_pem = private_key.public_key().public_bytes(
+        encoding=serialization.Encoding.PEM,
+        format=serialization.PublicFormat.SubjectPublicKeyInfo,
+    ).decode('ascii')
+    return private_key_pem, public_key_pem
+
+
 def _make_body(mapping, source_key='source', preserve: bool = True):
     """
     Given a dict mapping Keboola tables to aliases, construct the body of
@@ -44,6 +77,7 @@ def __init__(self, root_url, token):
             token (:obj:`str`): A storage API key.
         """
         super().__init__(root_url, 'workspaces', token)
+        self._default_backend = _DEFAULT_BACKEND_UNRESOLVED
 
     def list(self):
         """
@@ -61,8 +95,9 @@ def detail(self, workspace_id):
         """
         Retrieves information about a given workspace.
 
-        Note that the password to the workspace can only be retrieved when the
-        workspace is created.
+        Note that the workspace credentials (password or private key,
+        depending on the login type) are only available when the workspace
+        is created and cannot be retrieved later.
 
         Args:
             workspace_id (int or str): The id of the workspace.
@@ -77,24 +112,78 @@ def create(self, backend=None, timeout=None, login_type=None, public_key=None, r
         """
         Create a new Workspace and return the credentials.
 
+        On the snowflake backend, an omitted (or 'default') login_type would
+        create a deprecated password-based workspace. The client therefore
+        defaults to the 'snowflake-service-keypair' login type: when no
+        public_key is supplied, an RSA key pair is generated locally, the
+        public key is sent to the API and the private key is returned in
+        response['connection']['privateKey'] (it never leaves the client
+        otherwise and cannot be retrieved later). To get the deprecated
+        password-based workspace, pass login_type='snowflake-legacy-service'
+        explicitly.
+
         Args:
             backend (:obj:`str`): The type of engine for the workspace.
                 'redshift', 'snowflake' or 'synapse'. Defaults to the project's default backend.
             timeout (int): The timeout, in seconds, for SQL statements.
                 Only supported by snowflake backends.
+            login_type (:obj:`str`): The login type of the workspace, e.g.
+                'snowflake-service-keypair', 'snowflake-person-keypair',
+                'snowflake-legacy-service' or 'none'. Defaults to
+                'snowflake-service-keypair' on snowflake, otherwise to the
+                backend's default.
+            public_key (:obj:`str`): PEM-encoded RSA public key to use with
+                key-pair login types. When omitted for the default snowflake
+                key-pair login, a key pair is generated locally.
+            read_all_objects (bool): Grant the workspace read-only access to
+                all project data.
 
         Raises:
             requests.HTTPError: If the API request fails.
         """
+        private_key = None
+        effective_backend = backend or self._get_default_backend()
+        if effective_backend is None and login_type is not None:
+            raise ValueError(
+                "Cannot resolve the project's default backend from the token; "
+                "pass backend explicitly when using login_type."
+            )
+        if effective_backend == BACKEND_SNOWFLAKE:
+            if login_type in (None, LOGIN_TYPE_DEFAULT):
+                login_type = LOGIN_TYPE_SNOWFLAKE_SERVICE_KEYPAIR
+            if login_type == LOGIN_TYPE_SNOWFLAKE_SERVICE_KEYPAIR and public_key is None:
+                private_key, public_key = _generate_rsa_key_pair()
+        if login_type is not None:
+            # the API rejects loginType without an explicit backend
+            backend = effective_backend
+
         body = {
-            'backend': backend,
-            'statementTimeoutSeconds': timeout,
-            'loginType': login_type,
-            'publicKey': public_key,
-            'readOnlyStorageAccess': str(read_all_objects).lower()  # convert bool to lowercase true or false
+            k: v for k, v in {
+                'backend': backend,
+                'statementTimeoutSeconds': timeout,
+                'loginType': login_type,
+                'publicKey': public_key,
+                'readOnlyStorageAccess': str(read_all_objects).lower()  # convert bool to lowercase true or false
+            }.items()
+            if v is not None
         }
 
-        return self._post(self.base_url, data=body)
+        response = self._post(self.base_url, data=body)
+        if private_key is not None:
+            response.setdefault('connection', {})['privateKey'] = private_key
+        return response
+
+    def _get_default_backend(self):
+        """
+        Resolve the project's default backend from the token detail.
+
+        The value is cached on the instance - a project's default backend is
+        effectively immutable for the client's lifetime.
+        """
+        if self._default_backend is _DEFAULT_BACKEND_UNRESOLVED:
+            token_info = Tokens(self.root_url, self.token).verify()
+            self._default_backend = (token_info.get('owner') or {}).get('defaultBackend')
+        return self._default_backend
 
     def delete(self, workspace_id):
         """
@@ -116,6 +205,11 @@ def reset_password(self, workspace_id):
         """
         Generate a new password for the workspace.
 
+        Only supported for password-based login types (e.g. the deprecated
+        'snowflake-legacy-service'). For key-pair workspaces rotate the
+        credentials with set_public_key() using a freshly generated key pair
+        instead.
+
         Args:
             workspace_id (int or str): The id of the workspace for which the
                 password should be reset.

pyproject.toml

@@ -21,6 +21,7 @@ classifiers = [
 dependencies = [
     "boto3",
     "azure-storage-blob",
+    "cryptography",
     "urllib3<2.0.0",  # Frozen until fixed: https://github.com/boto/botocore/issues/2926
     # Dev dependencies
     "requests",

tests/mocks/test_workspaces.py

@@ -1,14 +1,19 @@
 """
-Asses basic functionality of the Workspace endpoint.
+Assess basic functionality of the Workspace endpoint.
 """
+import copy
 import unittest
+from urllib.parse import parse_qs
+
 import responses
 from requests import HTTPError
 
 from kbcstorage.workspaces import Workspaces
 
+from .token_responses import verify_token_response
 from .workspace_responses import (list_response, detail_response,
                                   load_tables_response, create_response,
+                                  keypair_create_response,
                                   reset_password_response)
 
 
@@ -91,18 +96,191 @@ def test_detail_inexsitent_workspace(self):
     @responses.activate
     def test_create(self):
         """
-        Workspace endpoint mock creates new workspace
+        Workspace endpoint mock creates new workspace. With no backend given,
+        the project default backend (snowflake) is resolved from the token and
+        a key-pair workspace is created instead of a password one.
         """
+        responses.add(
+            responses.Response(
+                method='GET',
+                url='https://connection.keboola.com/v2/storage/tokens/verify',
+                json=verify_token_response
+            )
+        )
         responses.add(
             responses.Response(
                 method='POST',
                 url='https://connection.keboola.com/v2/storage/workspaces',
-                json=create_response
+                json=keypair_create_response
             )
         )
         created_detail = self.ws.create()
+        request_body = parse_qs(responses.calls[1].request.body, keep_blank_values=True)
+        assert request_body['backend'] == ['snowflake']
+        assert request_body['loginType'] == ['snowflake-service-keypair']
+        assert 'BEGIN PUBLIC KEY' in request_body['publicKey'][0]
+        assert 'BEGIN PRIVATE KEY' in created_detail['connection']['privateKey']
+
+    @responses.activate
+    def test_create_snowflake_defaults_to_keypair(self):
+        """
+        With an explicit snowflake backend no token verify is needed and the
+        login type defaults to snowflake-service-keypair with a locally
+        generated key pair.
+        """
+        responses.add(
+            responses.Response(
+                method='POST',
+                url='https://connection.keboola.com/v2/storage/workspaces',
+                json=keypair_create_response
+            )
+        )
+        created_detail = self.ws.create(backend='snowflake')
+        assert len(responses.calls) == 1
+        request_body = parse_qs(responses.calls[0].request.body, keep_blank_values=True)
+        assert request_body['loginType'] == ['snowflake-service-keypair']
+        assert 'BEGIN PUBLIC KEY' in request_body['publicKey'][0]
+        assert 'BEGIN PRIVATE KEY' in created_detail['connection']['privateKey']
+
+    @responses.activate
+    def test_create_snowflake_with_own_public_key(self):
+        """
+        A caller-supplied public key is passed through and no private key is
+        injected into the response.
+        """
+        responses.add(
+            responses.Response(
+                method='POST',
+                url='https://connection.keboola.com/v2/storage/workspaces',
+                json=keypair_create_response
+            )
+        )
+        created_detail = self.ws.create(backend='snowflake', public_key='my-public-key')
+        request_body = parse_qs(responses.calls[0].request.body, keep_blank_values=True)
+        assert request_body['loginType'] == ['snowflake-service-keypair']
+        assert request_body['publicKey'] == ['my-public-key']
+        assert 'privateKey' not in created_detail['connection']
+
+    @responses.activate
+    def test_create_snowflake_explicit_legacy_login_type(self):
+        """
+        An explicitly requested password login type is passed through
+        unchanged and no key pair is generated.
+        """
+        responses.add(
+            responses.Response(
+                method='POST',
+                url='https://connection.keboola.com/v2/storage/workspaces',
+                json=create_response
+            )
+        )
+        created_detail = self.ws.create(backend='snowflake', login_type='snowflake-legacy-service')
+        request_body = parse_qs(responses.calls[0].request.body, keep_blank_values=True)
+        assert request_body['loginType'] == ['snowflake-legacy-service']
+        assert 'publicKey' not in request_body
         assert created_detail['connection']['password'] == 'abc'
 
+    @responses.activate
+    def test_create_non_snowflake_default_backend_unchanged(self):
+        """
+        When the project default backend resolved from the token is not
+        snowflake, no login type, public key or backend is sent.
+        """
+        bigquery_verify_response = copy.deepcopy(verify_token_response)
+        bigquery_verify_response['owner']['defaultBackend'] = 'bigquery'
+        responses.add(
+            responses.Response(
+                method='GET',
+                url='https://connection.keboola.com/v2/storage/tokens/verify',
+                json=bigquery_verify_response
+            )
+        )
+        responses.add(
+            responses.Response(
+                method='POST',
+                url='https://connection.keboola.com/v2/storage/workspaces',
+                json=create_response
+            )
+        )
+        self.ws.create()
+        request_body = parse_qs(responses.calls[1].request.body, keep_blank_values=True)
+        assert 'backend' not in request_body
+        assert 'loginType' not in request_body
+        assert 'publicKey' not in request_body
+
+    @responses.activate
+    def test_create_caches_default_backend(self):
+        """
+        The default backend is resolved via the token verify call only once
+        per endpoint instance.
+        """
+        responses.add(
+            responses.Response(
+                method='GET',
+                url='https://connection.keboola.com/v2/storage/tokens/verify',
+                json=verify_token_response
+            )
+        )
+        responses.add(
+            responses.Response(
+                method='POST',
+                url='https://connection.keboola.com/v2/storage/workspaces',
+                json=keypair_create_response
+            )
+        )
+        responses.add(
+            responses.Response(
+                method='POST',
+                url='https://connection.keboola.com/v2/storage/workspaces',
+                json=keypair_create_response
+            )
+        )
+        self.ws.create()
+        self.ws.create()
+        verify_calls = [c for c in responses.calls if c.request.method == 'GET']
+        assert len(verify_calls) == 1
+
+    @responses.activate
+    def test_create_login_type_without_resolvable_backend_raises(self):
+        """
+        The API rejects loginType without an explicit backend, so the client
+        raises a clear error when the default backend cannot be resolved. The
+        missing value is cached too - verify is not re-queried on retry.
+        """
+        no_backend_verify_response = copy.deepcopy(verify_token_response)
+        del no_backend_verify_response['owner']['defaultBackend']
+        responses.add(
+            responses.Response(
+                method='GET',
+                url='https://connection.keboola.com/v2/storage/tokens/verify',
+                json=no_backend_verify_response
+            )
+        )
+        with self.assertRaises(ValueError):
+            self.ws.create(login_type='none')
+        with self.assertRaises(ValueError):
+            self.ws.create(login_type='none')
+        assert len(responses.calls) == 1
+
+    @responses.activate
+    def test_create_non_snowflake_backend_unchanged(self):
+        """
+        Non-snowflake backends keep the original behavior - no login type or
+        public key is sent.
+        """
+        responses.add(
+            responses.Response(
+                method='POST',
+                url='https://connection.keboola.com/v2/storage/workspaces',
+                json=create_response
+            )
+        )
+        self.ws.create(backend='bigquery')
+        request_body = parse_qs(responses.calls[0].request.body, keep_blank_values=True)
+        assert request_body['backend'] == ['bigquery']
+        assert 'loginType' not in request_body
+        assert 'publicKey' not in request_body
+
     @responses.activate
     def test_delete(self):
         """