Snowflake: Support passing OAuth scope parameter in client_credentials flow

Author: vishalup29

providers/snowflake/docs/connections/snowflake.rst

@@ -60,6 +60,8 @@ Extra (optional)
     * ``authenticator``: To connect using OAuth set this parameter ``oauth``.
     * ``token_endpoint``: Specify token endpoint for external OAuth provider.
     * ``grant_type``: Specify grant type for OAuth authentication. Currently supported: ``refresh_token`` (default), ``client_credentials``.
+    * ``oauth_scope``: Optional OAuth scope to include when using the ``client_credentials`` grant type.
+      Some identity providers (e.g., Okta, Auth0) require a scope when issuing client credentials tokens.
     * ``refresh_token``: Specify refresh_token for OAuth connection.
     * ``azure_conn_id``: Azure Connection ID to be used for retrieving the OAuth token using Azure Entra authentication. Login and Password fields aren't required when using this method. Scope for the Azure OAuth token can be set in the config option ``azure_oauth_scope`` under the section ``[snowflake]``. Requires `apache-airflow-providers-microsoft-azure>=12.8.0`.
     * ``private_key_file``: Specify the path to the private key file.

providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake.py

@@ -224,11 +224,15 @@ def get_oauth_token(
         }
 
         if grant_type == "refresh_token":
-            data |= {
-                "refresh_token": conn_config["refresh_token"],
-            }
+            # Older provider versions may not supply refresh_token; avoid KeyError
+            refresh_token = conn_config.get("refresh_token")
+            if not refresh_token:
+                raise AirflowException("grant_type=refresh_token requires `refresh_token` in extra")
+            data["refresh_token"] = refresh_token
         elif grant_type == "client_credentials":
-            pass  # no setup necessary for client credentials grant.
+            oauth_scope = conn_config.get("oauth_scope")
+            if oauth_scope:
+                data["scope"] = oauth_scope
         else:
             raise ValueError(f"Unknown grant_type: {grant_type}")
 
@@ -379,6 +383,10 @@ def _get_conn_params(self) -> dict[str, str | None]:
             conn_config.pop("password", None)
 
         refresh_token = self._get_field(extra_dict, "refresh_token") or ""
+        oauth_scope = self._get_field(extra_dict, "oauth_scope")
+        if oauth_scope:
+            conn_config["oauth_scope"] = oauth_scope
+
         if refresh_token:
             conn_config["refresh_token"] = refresh_token
             conn_config["authenticator"] = "oauth"

providers/snowflake/tests/unit/snowflake/hooks/test_snowflake.py

@@ -1133,3 +1133,79 @@ def __init__(self, sdk_client, conn_id="azure_default"):
 
         # Check AzureBaseHook initialization
         mock_connection_class.get.assert_called_once_with(azure_conn_id)
+
+    @mock.patch("airflow.providers.snowflake.hooks.snowflake.HTTPBasicAuth")
+    @mock.patch("requests.post")
+    @mock.patch(
+        "airflow.providers.snowflake.hooks.snowflake.SnowflakeHook._get_conn_params",
+        new_callable=PropertyMock,
+    )
+    def test_get_oauth_token_client_credentials_without_scope(
+        self, mock_conn_params, requests_post,mock_auth
+    ):
+        """Test client_credentials flow without oauth_scope."""
+        mock_conn_params.return_value = {
+            **CONN_PARAMS_OAUTH_BASE,
+            "grant_type": "client_credentials",
+        }
+
+        basic_auth = {"Authorization": "Basic test"}
+        mock_auth.return_value = basic_auth
+
+        # Response mock
+        requests_post.return_value.status_code = 200
+        requests_post.return_value.json = lambda: {"access_token": "token123"}
+
+        hook = SnowflakeHook(snowflake_conn_id="test_conn")
+        hook.get_oauth_token(
+            conn_config=mock_conn_params.return_value,
+            grant_type="client_credentials",
+        )
+
+        requests_post.assert_called_once_with(
+            f"https://{CONN_PARAMS_OAUTH_BASE['account']}.snowflakecomputing.com/oauth/token-request",
+            data={
+                "grant_type": "client_credentials",
+                "redirect_uri": "https://localhost.com",
+            },
+            headers={"Content-Type": "application/x-www-form-urlencoded"},
+            auth=basic_auth,
+        )
+
+    @mock.patch("airflow.providers.snowflake.hooks.snowflake.HTTPBasicAuth")
+    @mock.patch("requests.post")
+    @mock.patch(
+        "airflow.providers.snowflake.hooks.snowflake.SnowflakeHook._get_conn_params",
+        new_callable=PropertyMock,
+    )
+    def test_get_oauth_token_client_credentials_with_scope(self, mock_conn_params, requests_post, mock_auth):
+        """Test client_credentials flow WITH oauth_scope."""
+        mock_conn_params.return_value = {
+            **CONN_PARAMS_OAUTH_BASE,
+            "grant_type": "client_credentials",
+            "oauth_scope": "custom_scope",
+        }
+
+        basic_auth = {"Authorization": "Basic test"}
+        mock_auth.return_value = basic_auth
+
+        # Response mock
+        requests_post.return_value.status_code = 200
+        requests_post.return_value.json = lambda: {"access_token": "token123"}
+
+        hook = SnowflakeHook(snowflake_conn_id="test_conn")
+        hook.get_oauth_token(
+            conn_config=mock_conn_params.return_value,
+            grant_type="client_credentials",
+        )
+
+        requests_post.assert_called_once_with(
+            f"https://{CONN_PARAMS_OAUTH_BASE['account']}.snowflakecomputing.com/oauth/token-request",
+            data={
+                "grant_type": "client_credentials",
+                "redirect_uri": "https://localhost.com",
+                "scope": "custom_scope",
+            },
+            headers={"Content-Type": "application/x-www-form-urlencoded"},
+            auth=basic_auth,
+        )