fix: Support role and warehouse as query parameters (#15)

Signed-off-by: Edgar Ramírez Mondragón <edgarrm358@gmail.com>

Author: edgarrmondragon

README.md

@@ -29,14 +29,27 @@ pipx inject meltano git+https://github.com/meltano/meltano-state-backend-snowfla
 
 ## Configuration
 
-To store state in Snowflake, set the `state_backend.uri` setting to `snowflake://<user>:<password>@<account>/<database>/<schema>`.
+To store state in Snowflake, set the `state_backend.uri` setting to a [Snowflake SQLAlchemy][snowflake-sqlalchemy]-style URI:
+
+```
+snowflake://<user>:<password>@<account>/<database>/<schema>?warehouse=<warehouse>&role=<role>
+```
 
 State will be stored in two tables that Meltano will create automatically:
 
 - `meltano_state` - Stores the actual state data
 - `meltano_state_locks` - Manages concurrency locks
 
-To authenticate to Snowflake, you'll need to provide:
+All connection parameters can be provided in the URI (including as query parameters), as individual Meltano settings, or a mix of both. Explicit settings take precedence over URI values.
+
+Using a single URI with query parameters:
+
+```yaml
+state_backend:
+  uri: snowflake://my_user:my_password@my_account/my_database/my_schema?warehouse=my_warehouse&role=my_role
+```
+
+Using a URI with separate settings for warehouse and role:
 
 ```yaml
 state_backend:
@@ -46,11 +59,11 @@ state_backend:
     role: my_role           # Optional: The role to use for the connection
 ```
 
-Alternatively, you can provide credentials via individual settings:
+Using individual settings for everything:
 
 ```yaml
 state_backend:
-  uri: snowflake://my_account/my_database/my_schema
+  uri: snowflake://my_account
   snowflake:
     account: my_account
     user: my_user
@@ -83,10 +96,11 @@ Example using environment variables:
 
 ```bash
 export MELTANO_STATE_BACKEND_SNOWFLAKE_PASSWORD='my_secure_password'
-meltano config set meltano state_backend.uri 'snowflake://my_user@my_account/my_database'
-meltano config set meltano state_backend.snowflake.warehouse 'my_warehouse'
+meltano config set meltano state_backend.uri 'snowflake://my_user@my_account/my_database?warehouse=my_warehouse'
 ```
 
+Passwords containing special characters (e.g. `@`, `%`) must be URL-encoded when included in the URI. For example, `p@ss` becomes `p%40ss`.
+
 ## Development
 
 ### Setup
@@ -115,5 +129,6 @@ gh release create v<new-version>
 [meltano]: https://meltano.com
 [pipx]: https://github.com/pypa/pipx
 [snowflake]: https://www.snowflake.com/
+[snowflake-sqlalchemy]: https://github.com/snowflakedb/snowflake-sqlalchemy
 [state-backend]: https://docs.meltano.com/concepts/state_backends
 [uv]: https://docs.astral.sh/uv

src/meltano_state_backend_snowflake/backend.py

@@ -7,7 +7,7 @@
 from contextlib import contextmanager
 from functools import cached_property
 from time import sleep
-from urllib.parse import urlparse
+from urllib.parse import parse_qs, urlparse
 
 import snowflake.connector
 from meltano.core.error import MeltanoError
@@ -123,6 +123,7 @@ def __init__(
         super().__init__(**kwargs)
         self.uri = uri
         parsed = urlparse(uri)
+        query_params = parse_qs(parsed.query)
 
         # Extract connection details from URI and parameters
         self.account = account or parsed.hostname
@@ -140,7 +141,7 @@ def __init__(
             msg = "Snowflake password is required"
             raise MissingStateBackendSettingsError(msg)
 
-        self.warehouse = warehouse
+        self.warehouse = warehouse or query_params.get("warehouse", [None])[0]
         if not self.warehouse:
             msg = "Snowflake warehouse is required"
             raise MissingStateBackendSettingsError(msg)
@@ -153,7 +154,7 @@ def __init__(
             raise MissingStateBackendSettingsError(msg)
 
         self.schema = schema or (path_parts[1] if len(path_parts) > 1 else "PUBLIC")
-        self.role = role
+        self.role = role or query_params.get("role", [None])[0]
 
         self._ensure_tables()
 

tests/test_backend.py

@@ -379,6 +379,66 @@ def test_acquire_lock_retry(
     assert mock_cursor.execute.call_count >= 2
 
 
+@pytest.mark.usefixtures("mock_connection")
+class TestURIQueryParams:
+    """Tests for URI query parameter parsing."""
+
+    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"
+
+    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
+
+
 def test_missing_account_validation() -> None:
     """Test missing account validation."""
     with pytest.raises(