copy code over from target-snowflake & add broser authentication

MeltanoLabs · Sep 20, 2024 · b90db8a · b90db8a
1 parent 86a73b2
commit b90db8a
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 5 deletions.
diff --git a/tap_snowflake/client.py b/tap_snowflake/client.py
@@ -54,6 +54,14 @@ def patched_conform(
 singer_sdk.helpers._typing._conform_primitive_property = patched_conform
 
 
+class SnowflakeAuthMethod(Enum):
+    """Supported methods to authenticate to snowflake"""
+
+    BROWSER = 1
+    PASSWORD = 2
+    KEY_PAIR = 3
+
+
 class ProfileStats(Enum):
     """Profile Statistics Enum."""
 
@@ -111,14 +119,23 @@ def get_private_key(self):
         )
 
     @cached_property
-    def auth_method(self):
+    def auth_method(self) -> SnowflakeAuthMethod:
         """Validate & return the authentication method based on config."""
+        if self.config.get("use_browser_authentication"):
+            return SnowflakeAuthMethod.BROWSER
+
         valid_auth_methods = {"private_key", "private_key_path", "password"}
         config_auth_methods = [x for x in self.config if x in valid_auth_methods]
         if len(config_auth_methods) != 1:
-            msg = f"One of {valid_auth_methods} must be specified"
+            msg = (
+                "Neither password nor private key was provided for "
+                "authentication. For password-less browser authentication via SSO, "
+                "set use_browser_authentication config option to True."
+            )
             raise ConfigValidationError(msg)
-        return config_auth_methods[0]
+        if config_auth_methods[0] in ["private_key", "private_key_path"]:
+            return SnowflakeAuthMethod.KEY_PAIR
+        return SnowflakeAuthMethod.PASSWORD
 
     def get_sqlalchemy_url(self, config: dict) -> str:
         """Concatenate a SQLAlchemy URL for use in connecting to the source."""
@@ -127,7 +144,9 @@ def get_sqlalchemy_url(self, config: dict) -> str:
             "user": config["user"],
         }
 
-        if self.auth_method == "password":
+        if self.auth_method == SnowflakeAuthMethod.BROWSER:
+            params["authenticator"] = "externalbrowser"
+        elif self.auth_method == SnowflakeAuthMethod.PASSWORD:
             params["password"] = config["password"]
 
         for option in ["database", "schema", "warehouse", "role"]:
@@ -143,7 +162,7 @@ def create_engine(self) -> sqlalchemy.engine.Engine:
             A SQLAlchemy engine.
         """
         connect_args = {}
-        if self.auth_method in ["private_key", "private_key_path"]:
+        if self.auth_method == SnowflakeAuthMethod.KEY_PAIR:
             connect_args["private_key"] = self.get_private_key()
         return sqlalchemy.create_engine(
             self.sqlalchemy_url,

diff --git a/tap_snowflake/tap.py b/tap_snowflake/tap.py
@@ -59,6 +59,16 @@ class TapSnowflake(SQLTap):
             secret=True,
             description="The passprhase used to protect the private key",
         ),
+        th.Property(
+            "use_browser_authentication",
+            th.BooleanType,
+            required=False,
+            default=False,
+            description=(
+                "If authentication should be done using SSO (via external browser). "
+                "See See SSO browser authentication."
+            )
+        ),
         th.Property(
             "account",
             th.StringType,