feat: Support Snowflake private key as a base64-encoded query parameter (#16)

Signed-off-by: Edgar Ramírez Mondragón <edgarrm358@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: edgarrmondragon

README.md

@@ -78,18 +78,49 @@ state_backend:
 
 - **account**: Your Snowflake account identifier (e.g., `myorg-account123`)
 - **user**: The username for authentication
-- **password**: The password for authentication
+- **password**: The password for authentication (required unless using key pair authentication)
 - **warehouse**: The compute warehouse to use (required)
 - **database**: The database where state will be stored
 - **schema**: The schema where state tables will be created (defaults to PUBLIC)
 - **role**: Optional role to use for the connection
+- **private_key_base64**: Optional base64-encoded DER private key for key pair authentication
+
+#### Key Pair Authentication
+
+Instead of password-based authentication, you can use [Snowflake key pair authentication][snowflake-keypair]. Provide the private key as a base64-encoded DER-format string:
+
+```yaml
+state_backend:
+  uri: snowflake://my_user@my_account/my_database?warehouse=my_warehouse
+  snowflake:
+    private_key_base64: MIIEvgIBADANBg...  # base64-encoded DER private key
+```
+
+The private key can also be passed as a URI query parameter:
+
+```
+snowflake://my_user@my_account/my_database?warehouse=my_warehouse&private_key_base64=MIIEvgIBADANBg...
+```
+
+Or via an environment variable:
+
+```bash
+export MELTANO_STATE_BACKEND_SNOWFLAKE_PRIVATE_KEY_BASE64='MIIEvgIBADANBg...'
+```
+
+To generate the base64-encoded DER key from a PEM private key file:
+
+```bash
+openssl pkcs8 -topk8 -inform PEM -outform DER -in rsa_key.pem -nocrypt | base64
+```
+
+When using key pair authentication, no password is required.
 
 #### Security Considerations
 
 When storing credentials:
 
 - Use environment variables for sensitive values in production
-- Consider using Snowflake key-pair authentication (future enhancement)
 - Ensure the user has CREATE TABLE, INSERT, UPDATE, DELETE, and SELECT privileges
 
 Example using environment variables:
@@ -129,6 +160,7 @@ gh release create v<new-version>
 [meltano]: https://meltano.com
 [pipx]: https://github.com/pypa/pipx
 [snowflake]: https://www.snowflake.com/
+[snowflake-keypair]: https://docs.snowflake.com/en/user-guide/key-pair-auth
 [snowflake-sqlalchemy]: https://github.com/snowflakedb/snowflake-sqlalchemy
 [state-backend]: https://docs.meltano.com/concepts/state_backends
 [uv]: https://docs.astral.sh/uv

fixtures/project/meltano.yml

@@ -1,4 +1,4 @@
-version: 1
+requires_meltano: ">=3.7"
 default_environment: dev
 project_id: 1ec76bd8-4499-4bad-a974-b27225d75f12
 send_anonymous_usage_stats: false

pyproject.toml

@@ -41,6 +41,7 @@ snowflake_role = "meltano_state_backend_snowflake.backend:SNOWFLAKE_ROLE"
 snowflake_schema = "meltano_state_backend_snowflake.backend:SNOWFLAKE_SCHEMA"
 snowflake_user = "meltano_state_backend_snowflake.backend:SNOWFLAKE_USER"
 snowflake_warehouse = "meltano_state_backend_snowflake.backend:SNOWFLAKE_WAREHOUSE"
+snowflake_private_key_base64 = "meltano_state_backend_snowflake.backend:SNOWFLAKE_PRIVATE_KEY_BASE64"
 
 [project.entry-points."meltano.state_backends"]
 snowflake = "meltano_state_backend_snowflake.backend:SnowflakeStateStoreManager"

src/meltano_state_backend_snowflake/backend.py

@@ -2,11 +2,12 @@
 
 from __future__ import annotations
 
+import base64
 import json
-import typing as t
 from contextlib import contextmanager
 from functools import cached_property
 from time import sleep
+from typing import TYPE_CHECKING, Any
 from urllib.parse import parse_qs, urlparse
 
 import snowflake.connector
@@ -19,7 +20,7 @@
     StateStoreManager,
 )
 
-if t.TYPE_CHECKING:
+if TYPE_CHECKING:
     from collections.abc import Generator, Iterable
 
 
@@ -85,6 +86,15 @@ class SnowflakeStateBackendError(MeltanoError):
     env_specific=True,
 )
 
+SNOWFLAKE_PRIVATE_KEY_BASE64 = SettingDefinition(
+    name="state_backend.snowflake.private_key_base64",
+    label="Snowflake Private Key (Base64-encoded DER)",
+    description="Snowflake private key to use, in base64-encoded DER format",
+    kind=SettingKind.STRING,  # ty: ignore[invalid-argument-type]
+    sensitive=True,
+    env_specific=True,
+)
+
 
 class SnowflakeStateStoreManager(StateStoreManager):
     """State backend for Snowflake."""
@@ -104,7 +114,8 @@ def __init__(
         database: str | None = None,
         schema: str | None = None,
         role: str | None = None,
-        **kwargs: t.Any,
+        private_key_base64: str | None = None,
+        **kwargs: Any,
     ) -> None:
         """Initialize the SnowflakeStateStoreManager.
 
@@ -117,6 +128,7 @@ def __init__(
             database: Snowflake database name
             schema: Snowflake schema name (default: PUBLIC)
             role: Optional Snowflake role to use
+            private_key_base64: Optional Snowflake private key to use, in base64-encoded DER format
             kwargs: Additional keyword args to pass to parent
 
         """
@@ -137,8 +149,12 @@ def __init__(
             raise MissingStateBackendSettingsError(msg)
 
         self.password = password or parsed.password
-        if not self.password:
-            msg = "Snowflake password is required"
+        if (
+            not self.password
+            and not private_key_base64
+            and not query_params.get("private_key_base64")
+        ):
+            msg = "Snowflake password or private key is required"
             raise MissingStateBackendSettingsError(msg)
 
         self.warehouse = warehouse or query_params.get("warehouse", [None])[0]
@@ -156,8 +172,18 @@ def __init__(
         self.schema = schema or (path_parts[1] if len(path_parts) > 1 else "PUBLIC")
         self.role = role or query_params.get("role", [None])[0]
 
+        pkey_base64 = private_key_base64 or query_params.get("private_key_base64", [None])[0]
+        self.private_key = self._load_private_key(pkey_base64) if pkey_base64 else None
+
         self._ensure_tables()
 
+    def _load_private_key(self, pkey_base64: str) -> bytes:
+        # Restore '+' chars that parse_qs decodes as spaces
+        pkey_base64 = pkey_base64.replace(" ", "+")
+        # Add padding if stripped (e.g. from URL query params)
+        pkey_base64 += "=" * (-len(pkey_base64) % 4)
+        return base64.b64decode(pkey_base64)
+
     @cached_property
     def connection(self) -> snowflake.connector.SnowflakeConnection:
         """Get a Snowflake connection.
@@ -166,17 +192,21 @@ def connection(self) -> snowflake.connector.SnowflakeConnection:
             A Snowflake connection object.
 
         """
-        conn_params = {
+        conn_params: dict[str, Any] = {
             "account": self.account,
             "user": self.user,
-            "password": self.password,
             "warehouse": self.warehouse,
             "database": self.database,
             "schema": self.schema,
         }
         if self.role:
             conn_params["role"] = self.role
 
+        if self.private_key:
+            conn_params["private_key"] = self.private_key
+        elif self.password:  # pragma: no branch
+            conn_params["password"] = self.password
+
         return snowflake.connector.connect(**conn_params)
 
     def _ensure_tables(self) -> None:

tests/test_backend.py

@@ -1,11 +1,14 @@
 from __future__ import annotations
 
+import base64
 import shutil
 from decimal import Decimal
 from typing import TYPE_CHECKING
 from unittest import mock
 
 import pytest
+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric import rsa
 from meltano.core.project import Project
 from meltano.core.state_store import MeltanoState, state_store_manager_from_project_settings
 from meltano.core.state_store.base import (
@@ -31,6 +34,17 @@ def project(tmp_path: Path) -> Project:
     return Project.find(path.resolve())  # type: ignore[no-any-return]
 
 
+@pytest.fixture
+def pkey_base64() -> str:
+    dummy_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
+    dummy_private_key = dummy_key.private_bytes(
+        encoding=serialization.Encoding.DER,
+        format=serialization.PrivateFormat.PKCS8,
+        encryption_algorithm=serialization.NoEncryption(),
+    )
+    return base64.b64encode(dummy_private_key).decode("utf-8")
+
+
 def test_get_manager(project: Project) -> None:
     with mock.patch(
         "meltano_state_backend_snowflake.backend.SnowflakeStateStoreManager._ensure_tables",
@@ -379,64 +393,78 @@ def test_acquire_lock_retry(
     assert mock_cursor.execute.call_count >= 2
 
 
+# class TestURIQueryParams:
+#     """Tests for URI query parameter parsing."""
+
+
+@pytest.fixture
+def base_uri() -> str:
+    return "snowflake://myuser:mypass@myaccount/mydb/myschema?warehouse=mywh"
+
+
 @pytest.mark.usefixtures("mock_connection")
-class TestURIQueryParams:
-    """Tests for URI query parameter parsing."""
+def test_full_sqlalchemy_uri(base_uri: str) -> None:
+    """Test full SQLAlchemy-style URI with all params."""
+    manager = SnowflakeStateStoreManager(uri=f"{base_uri}&role=myrole")
+    assert manager.account == "myaccount"
+    assert manager.user == "myuser"
+    assert manager.password == "mypass"  # noqa: S105
+    assert manager.database == "mydb"
+    assert manager.schema == "myschema"
+    assert manager.warehouse == "mywh"
+    assert manager.role == "myrole"
 
-    def test_full_sqlalchemy_uri(self) -> None:
-        """Test full SQLAlchemy-style URI with all params."""
-        manager = SnowflakeStateStoreManager(
-            uri="snowflake://myuser:mypass@myaccount/mydb/myschema?warehouse=mywh&role=myrole",
-        )
-        assert manager.account == "myaccount"
-        assert manager.user == "myuser"
-        assert manager.password == "mypass"  # noqa: S105
-        assert manager.database == "mydb"
-        assert manager.schema == "myschema"
-        assert manager.warehouse == "mywh"
-        assert manager.role == "myrole"
-
-    def test_warehouse_from_query_param(self) -> None:
-        """Test warehouse extracted from query parameter."""
-        manager = SnowflakeStateStoreManager(
-            uri="snowflake://myuser:mypass@myaccount/mydb?warehouse=mywh",
-        )
-        assert manager.warehouse == "mywh"
-        assert manager.schema == "PUBLIC"
 
-    def test_role_from_query_param(self) -> None:
-        """Test role extracted from query parameter."""
-        manager = SnowflakeStateStoreManager(
-            uri="snowflake://myuser:mypass@myaccount/mydb?warehouse=mywh&role=analyst",
-        )
-        assert manager.role == "analyst"
+@pytest.mark.usefixtures("mock_connection")
+def test_warehouse_from_query_param(base_uri: str) -> None:
+    """Test warehouse extracted from query parameter."""
+    manager = SnowflakeStateStoreManager(uri=base_uri)
+    assert manager.warehouse == "mywh"
+    assert manager.schema == "myschema"
 
-    def test_explicit_settings_override_query_params(self) -> None:
-        """Test that explicit settings take precedence over query params."""
-        manager = SnowflakeStateStoreManager(
-            uri="snowflake://uri_user:uri_pass@uri_account/uri_db/uri_schema?warehouse=uri_wh&role=uri_role",
-            account="explicit_account",
-            user="explicit_user",
-            password="explicit_pass",  # noqa: S106
-            warehouse="explicit_wh",
-            database="explicit_db",
-            schema="explicit_schema",
-            role="explicit_role",
-        )
-        assert manager.account == "explicit_account"
-        assert manager.user == "explicit_user"
-        assert manager.password == "explicit_pass"  # noqa: S105
-        assert manager.warehouse == "explicit_wh"
-        assert manager.database == "explicit_db"
-        assert manager.schema == "explicit_schema"
-        assert manager.role == "explicit_role"
-
-    def test_role_defaults_to_none(self) -> None:
-        """Test role defaults to None when not provided."""
-        manager = SnowflakeStateStoreManager(
-            uri="snowflake://myuser:mypass@myaccount/mydb?warehouse=mywh",
-        )
-        assert manager.role is None
+
+@pytest.mark.usefixtures("mock_connection")
+def test_role_from_query_param(base_uri: str) -> None:
+    """Test role extracted from query parameter."""
+    manager = SnowflakeStateStoreManager(uri=f"{base_uri}&role=analyst")
+    assert manager.role == "analyst"
+
+
+@pytest.mark.usefixtures("mock_connection")
+def test_explicit_settings_override_query_params(base_uri: str) -> None:
+    """Test that explicit settings take precedence over query params."""
+    manager = SnowflakeStateStoreManager(
+        uri=f"{base_uri}&role=uri_role",
+        account="explicit_account",
+        user="explicit_user",
+        password="explicit_pass",  # noqa: S106
+        warehouse="explicit_wh",
+        database="explicit_db",
+        schema="explicit_schema",
+        role="explicit_role",
+    )
+    assert manager.account == "explicit_account"
+    assert manager.user == "explicit_user"
+    assert manager.password == "explicit_pass"  # noqa: S105
+    assert manager.warehouse == "explicit_wh"
+    assert manager.database == "explicit_db"
+    assert manager.schema == "explicit_schema"
+    assert manager.role == "explicit_role"
+
+
+@pytest.mark.usefixtures("mock_connection")
+def test_role_defaults_to_none(base_uri: str) -> None:
+    """Test role defaults to None when not provided."""
+    manager = SnowflakeStateStoreManager(uri=base_uri)
+    assert manager.role is None
+
+
+@pytest.mark.usefixtures("mock_connection")
+def test_private_key_from_query_param(base_uri: str, pkey_base64: str) -> None:
+    """Test private key extracted from query parameter."""
+    manager = SnowflakeStateStoreManager(uri=f"{base_uri}&private_key_base64={pkey_base64}")
+    assert manager.private_key is not None
+    assert manager.private_key == base64.b64decode(pkey_base64)
 
 
 def test_missing_account_validation() -> None:
@@ -470,10 +498,10 @@ def test_missing_user_validation() -> None:
 
 
 def test_missing_password_validation() -> None:
-    """Test missing password validation."""
+    """Test missing password validation when no private key is provided."""
     with pytest.raises(
         MissingStateBackendSettingsError,
-        match="Snowflake password is required",
+        match="Snowflake password or private key is required",
     ):
         SnowflakeStateStoreManager(
             uri="snowflake://user@account/db",  # No password in URI
@@ -484,6 +512,17 @@ def test_missing_password_validation() -> None:
         )
 
 
+@pytest.mark.usefixtures("mock_connection")
+def test_private_key_without_password(pkey_base64: str) -> None:
+    """Test that private key auth works without a password."""
+    manager = SnowflakeStateStoreManager(
+        uri="snowflake://myuser@myaccount/mydb?warehouse=mywh",
+        private_key_base64=pkey_base64,
+    )
+    assert manager.private_key == base64.b64decode(pkey_base64)
+    assert manager.password is None
+
+
 def test_missing_warehouse_validation() -> None:
     """Test missing warehouse validation."""
     with pytest.raises(