Merge pull request #738 from hkad98/key-pair-snowflake-auth

Add key pair auth for Snowflake support

Author: hkad98

docs/Dockerfile

@@ -1,12 +1,13 @@
 FROM node:20.11.1-bookworm-slim
 
-RUN npm install -g [email protected]
+RUN apt-get update && \
+    apt-get install -y git make golang-go curl && \
+    npm install -g [email protected] && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*
 
 COPY docs docs
 
-RUN apt-get update && \
-    apt-get install -y git make golang-go curl
-
 WORKDIR docs
 RUN npm install
 

docs/content/en/latest/data/data-source/_index.md

@@ -40,7 +40,7 @@ See [Connect Data](https://www.gooddata.com/docs/cloud/connect-data/) to learn h
 
 ## Example
 
-Since there are multiple data source types, here are examples, how to initialize each of them:
+Since there are multiple data source types, here are examples of how to initialize each of them:
 
 ### Postgres
 
@@ -79,6 +79,7 @@ CatalogDataSourceRedshift(
 ```
 ### Snowflake
 
+Using basic credentials (username + password):
 ```python
 CatalogDataSourceSnowflake(
     id=data_source_id,
@@ -95,6 +96,25 @@ CatalogDataSourceSnowflake(
     ),
 )
 ```
+
+Using key-pair authentication (username + private_key, optionally private_key_passphrase):
+```python
+CatalogDataSourceSnowflake(
+    id=data_source_id,
+    name=data_source_name,
+    db_specific_attributes=SnowflakeAttributes(
+        account=os.environ["SNOWFLAKE_ACCOUNT"],
+        warehouse=os.environ["SNOWFLAKE_WAREHOUSE"],
+        db_name=os.environ["SNOWFLAKE_DBNAME"]
+    ),
+    schema=os.environ["SNOWFLAKE_SCHEMA"],
+    credentials=KeyPairCredentials(
+        username=os.environ["SNOWFLAKE_USER"],
+        private_key=os.environ["SNOWFLAKE_PRIVATE_KEY"],
+    ),
+)
+```
+
 ### Vertica
 
 ```python

gooddata-sdk/gooddata_sdk/__init__.py

@@ -56,6 +56,7 @@
 from gooddata_sdk.catalog.entity import (
     AttrCatalogEntity,
     BasicCredentials,
+    KeyPairCredentials,
     TokenCredentialsFromEnvVar,
     TokenCredentialsFromFile,
 )

gooddata-sdk/gooddata_sdk/catalog/data_source/declarative_model/data_source.py

@@ -107,6 +107,8 @@ def to_test_request(
         self,
         password: Optional[str] = None,
         token: Optional[str] = None,
+        private_key: Optional[str] = None,
+        private_key_passphrase: Optional[str] = None,
     ) -> TestDefinitionRequest:
         kwargs: dict[str, Any] = {"schema": self.schema}
         if password is not None:
@@ -115,6 +117,10 @@ def to_test_request(
             kwargs["token"] = token
         if self.username is not None:
             kwargs["username"] = self.username
+        if private_key is not None:
+            kwargs["private_key"] = private_key
+        if private_key_passphrase is not None:
+            kwargs["private_key_passphrase"] = private_key
         return TestDefinitionRequest(type=self.type, url=self.url, **kwargs)
 
     @staticmethod
@@ -127,12 +133,22 @@ def data_source_folder(data_sources_folder: Path, data_source_id: str) -> Path:
         create_directory(data_source_folder)
         return data_source_folder
 
-    def to_api(self, password: Optional[str] = None, token: Optional[str] = None) -> DeclarativeDataSource:
+    def to_api(
+        self,
+        password: Optional[str] = None,
+        token: Optional[str] = None,
+        private_key: Optional[str] = None,
+        private_key_passphrase: Optional[str] = None,
+    ) -> DeclarativeDataSource:
         dictionary = self._get_snake_dict()
         if password is not None:
             dictionary["password"] = password
         if token is not None:
             dictionary["token"] = token
+        if private_key is not None:
+            dictionary["private_key"] = private_key
+        if private_key_passphrase is not None:
+            dictionary["private_key_passphrase"] = private_key_passphrase
         return self.client_class().from_dict(dictionary)
 
     def store_to_disk(self, data_sources_folder: Path) -> None:

gooddata-sdk/gooddata_sdk/catalog/data_source/entity_model/data_source.py

@@ -13,7 +13,13 @@
 from gooddata_api_client.model.json_api_data_source_patch_document import JsonApiDataSourcePatchDocument
 
 from gooddata_sdk.catalog.base import Base, value_in_allowed
-from gooddata_sdk.catalog.entity import BasicCredentials, Credentials, TokenCredentials, TokenCredentialsFromFile
+from gooddata_sdk.catalog.entity import (
+    BasicCredentials,
+    Credentials,
+    KeyPairCredentials,
+    TokenCredentials,
+    TokenCredentialsFromFile,
+)
 
 U = TypeVar("U", bound="CatalogDataSourceBase")
 
@@ -29,6 +35,7 @@ class CatalogDataSourceBase(Base):
         BasicCredentials,
         TokenCredentials,
         TokenCredentialsFromFile,
+        KeyPairCredentials,
     ]
     _DELIMITER: ClassVar[str] = "&"
     _ATTRIBUTES: ClassVar[List[str]] = [

gooddata-sdk/gooddata_sdk/catalog/entity.py

@@ -118,6 +118,8 @@ class Credentials(Base):
     TOKEN_KEY: ClassVar[str] = "token"
     USER_KEY: ClassVar[str] = "username"
     PASSWORD_KEY: ClassVar[str] = "password"
+    PRIVATE_KEY: ClassVar[str] = "private_key"
+    PRIVATE_KEY_PASSPHRASE: ClassVar[str] = "private_key_passphrase"
 
     def to_api_args(self) -> dict[str, Any]:
         return attr.asdict(self)
@@ -228,3 +230,24 @@ def from_api(cls, attributes: dict[str, Any]) -> BasicCredentials:
             # You have to fill it to keep it or update it
             password="",
         )
+
+
+@attr.s(auto_attribs=True, kw_only=True)
+class KeyPairCredentials(Credentials):
+    username: str
+    private_key: str = attr.field(repr=lambda value: "***")
+    private_key_passphrase: Optional[str] = attr.field(repr=lambda value: "***", default=None)
+
+    @classmethod
+    def is_part_of_api(cls, entity: dict[str, Any]) -> bool:
+        return cls.USER_KEY in entity and cls.PRIVATE_KEY in entity
+
+    @classmethod
+    def from_api(cls, attributes: dict[str, Any]) -> KeyPairCredentials:
+        # Credentials are not returned for security reasons
+        return cls(
+            username=attributes[cls.USER_KEY],
+            # Private key is not returned from API (security)
+            # You have to fill it to keep it or update it
+            private_key="",
+        )