Environment variable to force browser-based auth (#2538)

Co-authored-by: Patryk Czajka <[email protected]>

(cherry picked from commit 54906d6)

Author: HeroCC

src/snowflake/connector/auth/webbrowser.py

@@ -165,15 +165,26 @@ def prepare(
                 return
 
             print(
-                "Initiating login request with your identity provider. A "
-                "browser window should have opened for you to complete the "
-                "login. If you can't see it, check existing browser windows, "
-                "or your OS settings. Press CTRL+C to abort and try again..."
+                "Initiating login request with your identity provider. Press CTRL+C to abort and try again..."
             )
 
             logger.debug("step 2: open a browser")
             print(f"Going to open: {sso_url} to authenticate...")
-            if not self._webbrowser.open_new(sso_url):
+            browser_opened = self._webbrowser.open_new(sso_url)
+            if browser_opened:
+                print(
+                    "A browser window should have opened for you to complete the "
+                    "login. If you can't see it, check existing browser windows, "
+                    "or your OS settings."
+                )
+
+            if (
+                browser_opened
+                or os.getenv("SNOWFLAKE_AUTH_FORCE_SERVER", "False").lower() == "true"
+            ):
+                logger.debug("step 3: accept SAML token")
+                self._receive_saml_token(conn, socket_connection)
+            else:
                 print(
                     "We were unable to open a browser window for you, "
                     "please open the url above manually then paste the "
@@ -195,9 +206,6 @@ def prepare(
                         },
                     )
                     return
-            else:
-                logger.debug("step 3: accept SAML token")
-                self._receive_saml_token(conn, socket_connection)
         finally:
             socket_connection.close()
 

test/unit/test_auth_webbrowser.py

@@ -277,9 +277,7 @@ def test_auth_webbrowser_fail_webbrowser(
         )
     captured = capsys.readouterr()
     assert captured.out == (
-        "Initiating login request with your identity provider. A browser window "
-        "should have opened for you to complete the login. If you can't see it, "
-        "check existing browser windows, or your OS settings. Press CTRL+C to "
+        "Initiating login request with your identity provider. Press CTRL+C to "
         f"abort and try again...\nGoing to open: {REF_SSO_URL if disable_console_login else REF_CONSOLE_LOGIN_SSO_URL} to authenticate...\nWe were unable to open a browser window for "
         "you, please open the url above manually then paste the URL you "
         "are redirected to into the terminal.\n"
@@ -336,10 +334,10 @@ def test_auth_webbrowser_fail_webserver(_, capsys, disable_console_login):
         )
         captured = capsys.readouterr()
         assert captured.out == (
-            "Initiating login request with your identity provider. A browser window "
+            "Initiating login request with your identity provider. Press CTRL+C to "
+            f"abort and try again...\nGoing to open: {REF_SSO_URL if disable_console_login else REF_CONSOLE_LOGIN_SSO_URL} to authenticate...\nA browser window "
             "should have opened for you to complete the login. If you can't see it, "
-            "check existing browser windows, or your OS settings. Press CTRL+C to "
-            f"abort and try again...\nGoing to open: {REF_SSO_URL if disable_console_login else REF_CONSOLE_LOGIN_SSO_URL} to authenticate...\n"
+            "check existing browser windows, or your OS settings.\n"
         )
         assert rest._connection.errorhandler.called  # an error
         assert auth.assertion_content is None
@@ -751,6 +749,78 @@ def test_auth_webbrowser_socket_reuseport_option_not_set_with_no_flag(monkeypatc
         assert auth.assertion_content == ref_token
 
 
+@pytest.mark.parametrize("force_auth_server", [True, False])
+@patch("secrets.token_bytes", return_value=PROOF_KEY)
+def test_auth_webbrowser_force_auth_server(_, monkeypatch, force_auth_server):
+    """Authentication by WebBrowser with SNOWFLAKE_AUTH_FORCE_SERVER environment variable."""
+    ref_token = "MOCK_TOKEN"
+    rest = _init_rest(REF_SSO_URL, REF_PROOF_KEY, disable_console_login=True)
+
+    # Set environment variable
+    if force_auth_server:
+        monkeypatch.setenv("SNOWFLAKE_AUTH_FORCE_SERVER", "true")
+    else:
+        monkeypatch.delenv("SNOWFLAKE_AUTH_FORCE_SERVER", raising=False)
+
+    # mock socket
+    mock_socket_pkg = _init_socket(
+        recv_side_effect_func=recv_setup([successful_web_callback(ref_token)])
+    )
+
+    # mock webbrowser - simulate browser failing to open
+    mock_webbrowser = MagicMock()
+    mock_webbrowser.open_new.return_value = False
+
+    # Mock select.select to return socket client
+    with mock.patch(
+        "select.select", return_value=([mock_socket_pkg.return_value], [], [])
+    ):
+        auth = AuthByWebBrowser(
+            application=APPLICATION,
+            webbrowser_pkg=mock_webbrowser,
+            socket_pkg=mock_socket_pkg,
+        )
+
+        if force_auth_server:
+            # When SNOWFLAKE_AUTH_FORCE_SERVER is true, should continue with server flow even if browser fails
+            auth.prepare(
+                conn=rest._connection,
+                authenticator=AUTHENTICATOR,
+                service_name=SERVICE_NAME,
+                account=ACCOUNT,
+                user=USER,
+                password=PASSWORD,
+            )
+            assert not rest._connection.errorhandler.called  # no error
+            assert auth.assertion_content == ref_token
+            body = {"data": {}}
+            auth.update_body(body)
+            assert body["data"]["TOKEN"] == ref_token
+            assert body["data"]["AUTHENTICATOR"] == EXTERNAL_BROWSER_AUTHENTICATOR
+            assert body["data"]["PROOF_KEY"] == REF_PROOF_KEY
+        else:
+            # When SNOWFLAKE_AUTH_FORCE_SERVER is false/unset, should fall back to manual URL input
+            with patch(
+                "builtins.input",
+                return_value=f"http://example.com/sso?token={ref_token}",
+            ):
+                auth.prepare(
+                    conn=rest._connection,
+                    authenticator=AUTHENTICATOR,
+                    service_name=SERVICE_NAME,
+                    account=ACCOUNT,
+                    user=USER,
+                    password=PASSWORD,
+                )
+                assert not rest._connection.errorhandler.called  # no error
+                assert auth.assertion_content == ref_token
+                body = {"data": {}}
+                auth.update_body(body)
+                assert body["data"]["TOKEN"] == ref_token
+                assert body["data"]["AUTHENTICATOR"] == EXTERNAL_BROWSER_AUTHENTICATOR
+                assert body["data"]["PROOF_KEY"] == REF_PROOF_KEY
+
+
 @pytest.mark.parametrize("authenticator", ["EXTERNALBROWSER", "externalbrowser"])
 def test_externalbrowser_authenticator_is_case_insensitive(monkeypatch, authenticator):
     """Test that external browser authenticator is case insensitive."""