Add support for Snowflake authenticator (#644)

* add support for snowflake authenticator

* dbt only load default connection

Author: eakmanrq

sqlmesh/core/config/connection.py

@@ -4,7 +4,7 @@
 import sys
 import typing as t
 
-from pydantic import Field
+from pydantic import Field, root_validator
 
 from sqlmesh.core import engine_adapter
 from sqlmesh.core.config.base import BaseConfig
@@ -13,6 +13,7 @@
     http_headers_validator,
 )
 from sqlmesh.core.engine_adapter import EngineAdapter
+from sqlmesh.utils.errors import ConfigError
 
 if sys.version_info >= (3, 9):
     from typing import Annotated, Literal
@@ -89,31 +90,45 @@ class SnowflakeConnectionConfig(_ConnectionConfig):
     """Configuration for the Snowflake connection.
 
     Args:
+        account: The Snowflake account name.
         user: The Snowflake username.
         password: The Snowflake password.
-        account: The Snowflake account name.
         warehouse: The optional warehouse name.
         database: The optional database name.
         role: The optional role name.
         concurrent_tasks: The maximum number of tasks that can use this connection concurrently.
+        authenticator: The optional authenticator name. Defaults to username/password authentication ("snowflake").
+                       Options: https://github.com/snowflakedb/snowflake-connector-python/blob/e937591356c067a77f34a0a42328907fda792c23/src/snowflake/connector/network.py#L178-L183
     """
 
-    user: str
-    password: str
     account: str
+    user: t.Optional[str]
+    password: t.Optional[str]
     warehouse: t.Optional[str]
     database: t.Optional[str]
     role: t.Optional[str]
+    authenticator: t.Optional[str]
 
     concurrent_tasks: int = 4
 
     type_: Literal["snowflake"] = Field(alias="type", default="snowflake")
 
     _concurrent_tasks_validator = concurrent_tasks_validator
 
+    @root_validator()
+    def _validate_authenticator(
+        cls, fields: t.Dict[str, t.Optional[str]]
+    ) -> t.Dict[str, t.Optional[str]]:
+        auth = fields.get("authenticator")
+        user = fields.get("user")
+        password = fields.get("password")
+        if not auth and (not user or not password):
+            raise ConfigError("User and password must be provided if using default authentication")
+        return fields
+
     @property
     def _connection_kwargs_keys(self) -> t.Set[str]:
-        return {"user", "password", "account", "warehouse", "database", "role"}
+        return {"user", "password", "account", "warehouse", "database", "role", "authenticator"}
 
     @property
     def _engine_adapter(self) -> t.Type[EngineAdapter]:

sqlmesh/dbt/profile.py

@@ -101,4 +101,4 @@ def _read_profile(
         return (targets, default_target)
 
     def to_sqlmesh(self) -> t.Dict[str, ConnectionConfig]:
-        return {name: target.to_sqlmesh() for name, target in self.targets.items()}
+        return {self.default_target: self.targets[self.default_target].to_sqlmesh()}

tests/core/test_connection_config.py

@@ -0,0 +1,22 @@
+import pytest
+
+from sqlmesh.core.config.connection import SnowflakeConnectionConfig
+from sqlmesh.utils.errors import ConfigError
+
+
+def test_snowflake_auth():
+    # Authenticator and user/password is fine
+    SnowflakeConnectionConfig(
+        account="test", user="test", password="test", authenticator="externalbrowser"
+    )
+    # Auth with no user/password is fine
+    SnowflakeConnectionConfig(account="test", authenticator="externalbrowser")
+    # No auth and no user raises
+    with pytest.raises(ConfigError):
+        SnowflakeConnectionConfig(account="test", password="test")
+    # No auth and no password raises
+    with pytest.raises(ConfigError):
+        SnowflakeConnectionConfig(account="test", user="test")
+    # No auth and no user/password raises
+    with pytest.raises(ConfigError):
+        SnowflakeConnectionConfig(account="test")