From 35fb2d0c1d4e21435ed8d9e6f5a0c57095e5f11e Mon Sep 17 00:00:00 2001 From: Tianyou Gu Date: Tue, 5 Mar 2024 15:32:52 -0800 Subject: [PATCH] feat: update SnowflakeSqlApiHook to support OAuth --- .../providers/snowflake/hooks/snowflake.py | 9 +++ .../snowflake/hooks/snowflake_sql_api.py | 74 +++++++++++++++---- .../connections/snowflake.rst | 3 +- .../operators/snowflake.rst | 6 +- .../snowflake/hooks/test_snowflake_sql_api.py | 62 +++++++++++++++- 5 files changed, 134 insertions(+), 20 deletions(-) diff --git a/airflow/providers/snowflake/hooks/snowflake.py b/airflow/providers/snowflake/hooks/snowflake.py index c545f668fb3ab..f7e989665315e 100644 --- a/airflow/providers/snowflake/hooks/snowflake.py +++ b/airflow/providers/snowflake/hooks/snowflake.py @@ -256,6 +256,15 @@ def _get_conn_params(self) -> dict[str, str | None]: conn_config["private_key"] = pkb conn_config.pop("password", None) + refresh_token = self._get_field(extra_dict, "refresh_token") or "" + if refresh_token: + conn_config["refresh_token"] = refresh_token + conn_config["authenticator"] = "oauth" + conn_config["client_id"] = conn.login + conn_config["client_secret"] = conn.password + conn_config.pop("login", None) + conn_config.pop("password", None) + return conn_config def get_uri(self) -> str: diff --git a/airflow/providers/snowflake/hooks/snowflake_sql_api.py b/airflow/providers/snowflake/hooks/snowflake_sql_api.py index 9d564144d100e..018182302ee4d 100644 --- a/airflow/providers/snowflake/hooks/snowflake_sql_api.py +++ b/airflow/providers/snowflake/hooks/snowflake_sql_api.py @@ -25,6 +25,7 @@ import requests from cryptography.hazmat.backends import default_backend from cryptography.hazmat.primitives import serialization +from requests.auth import HTTPBasicAuth from airflow.exceptions import AirflowException from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook @@ -39,8 +40,9 @@ class SnowflakeSqlApiHook(SnowflakeHook): poll to check the status of the execution of a statement. Fetch query results asynchronously. This hook requires the snowflake_conn_id connection. This hooks mainly uses account, schema, database, - warehouse, private_key_file or private_key_content field must be setup in the connection. Other inputs - can be defined in the connection or hook instantiation. + warehouse, and an authentication mechanism from one of below: + 1. JWT Token generated from private_key_file or private_key_content. Other inputs can be defined in the connection or hook instantiation. + 2. OAuth Token generated from the refresh_token, client_id and client_secret specified in the connection :param snowflake_conn_id: Reference to :ref:`Snowflake connection id` @@ -81,6 +83,17 @@ def __init__( super().__init__(snowflake_conn_id, *args, **kwargs) self.private_key: Any = None + @property + def account_identifier(self) -> str: + """Returns snowflake account identifier.""" + conn_config = self._get_conn_params() + account_identifier = f"https://{conn_config['account']}" + + if conn_config["region"]: + account_identifier += f".{conn_config['region']}" + + return account_identifier + def get_private_key(self) -> None: """Get the private key from snowflake connection.""" conn = self.get_connection(self.snowflake_conn_id) @@ -137,10 +150,7 @@ def execute_query( conn_config = self._get_conn_params() req_id = uuid.uuid4() - url = ( - f"https://{conn_config['account']}.{conn_config['region']}" - f".snowflakecomputing.com/api/v2/statements" - ) + url = f"{self.account_identifier}.snowflakecomputing.com/api/v2/statements" params: dict[str, Any] | None = {"requestId": str(req_id), "async": True, "pageSize": 10} headers = self.get_headers() if bindings is None: @@ -175,12 +185,27 @@ def execute_query( return self.query_ids def get_headers(self) -> dict[str, Any]: - """Form JWT Token and header based on the private key, and connection details.""" + """Form auth headers based on either OAuth token or JWT token from private key.""" + conn_config = self._get_conn_params() + + # Use OAuth if refresh_token and client_id and client_secret are provided + if all( + [conn_config.get("refresh_token"), conn_config.get("client_id"), conn_config.get("client_secret")] + ): + oauth_token = self.get_oauth_token() + headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {oauth_token}", + "Accept": "application/json", + "User-Agent": "snowflakeSQLAPI/1.0", + "X-Snowflake-Authorization-Token-Type": "OAUTH", + } + return headers + + # Alternatively, get the JWT token from the connection details and the private key if not self.private_key: self.get_private_key() - conn_config = self._get_conn_params() - # Get the JWT token from the connection details and the private key token = JWTGenerator( conn_config["account"], # type: ignore[arg-type] conn_config["user"], # type: ignore[arg-type] @@ -198,20 +223,41 @@ def get_headers(self) -> dict[str, Any]: } return headers + def get_oauth_token(self) -> str: + """Generate temporary OAuth access token using refresh token in connection details.""" + conn_config = self._get_conn_params() + url = f"{self.account_identifier}.snowflakecomputing.com/oauth/token-request" + data = { + "grant_type": "refresh_token", + "refresh_token": conn_config["refresh_token"], + "redirect_uri": conn_config.get("redirect_uri", "https://localhost.com"), + } + response = requests.post( + url, + data=data, + headers={ + "Content-Type": "application/x-www-form-urlencoded", + }, + auth=HTTPBasicAuth(conn_config["client_id"], conn_config["client_secret"]), # type: ignore[arg-type] + ) + + try: + response.raise_for_status() + except requests.exceptions.HTTPError as e: # pragma: no cover + msg = f"Response: {e.response.content.decode()} Status Code: {e.response.status_code}" + raise AirflowException(msg) + return response.json()["access_token"] + def get_request_url_header_params(self, query_id: str) -> tuple[dict[str, Any], dict[str, Any], str]: """ Build the request header Url with account name identifier and query id from the connection params. :param query_id: statement handles query ids for the individual statements. """ - conn_config = self._get_conn_params() req_id = uuid.uuid4() header = self.get_headers() params = {"requestId": str(req_id)} - url = ( - f"https://{conn_config['account']}.{conn_config['region']}" - f".snowflakecomputing.com/api/v2/statements/{query_id}" - ) + url = f"{self.account_identifier}.snowflakecomputing.com/api/v2/statements/{query_id}" return header, params, url def check_query_output(self, query_ids: list[str]) -> None: diff --git a/docs/apache-airflow-providers-snowflake/connections/snowflake.rst b/docs/apache-airflow-providers-snowflake/connections/snowflake.rst index e4c2f39234647..8d4a3aad646ea 100644 --- a/docs/apache-airflow-providers-snowflake/connections/snowflake.rst +++ b/docs/apache-airflow-providers-snowflake/connections/snowflake.rst @@ -56,7 +56,8 @@ Extra (optional) * ``region``: Warehouse region. * ``warehouse``: Snowflake warehouse name. * ``role``: Snowflake role. - * ``authenticator``: To connect using OAuth set this parameter ``oath``. + * ``authenticator``: To connect using OAuth set this parameter ``oauth``. + * ``refresh_token``: Specify refresh_token for OAuth connection. * ``private_key_file``: Specify the path to the private key file. * ``private_key_content``: Specify the content of the private key file. * ``session_parameters``: Specify `session level parameters `_. diff --git a/docs/apache-airflow-providers-snowflake/operators/snowflake.rst b/docs/apache-airflow-providers-snowflake/operators/snowflake.rst index d3ffbec5a4cfe..5a94f5b87e07e 100644 --- a/docs/apache-airflow-providers-snowflake/operators/snowflake.rst +++ b/docs/apache-airflow-providers-snowflake/operators/snowflake.rst @@ -82,13 +82,13 @@ the connection metadata is structured as follows: * - Parameter - Input * - Login: string - - Snowflake user name + - Snowflake user name. If using `OAuth connection `__ this is the ``client_id`` * - Password: string - - Password for Snowflake user + - Password for Snowflake user. If using OAuth this is the ``client_secret`` * - Schema: string - Set schema to execute SQL operations on by default * - Extra: dictionary - - ``warehouse``, ``account``, ``database``, ``region``, ``role``, ``authenticator`` + - ``warehouse``, ``account``, ``database``, ``region``, ``role``, ``authenticator``, ``refresh_token``. If using OAuth must specify ``refresh_token`` (`obtained here `__) An example usage of the SnowflakeSqlApiHook is as follows: diff --git a/tests/providers/snowflake/hooks/test_snowflake_sql_api.py b/tests/providers/snowflake/hooks/test_snowflake_sql_api.py index 6fc8ad983f76f..bfd978755e6dd 100644 --- a/tests/providers/snowflake/hooks/test_snowflake_sql_api.py +++ b/tests/providers/snowflake/hooks/test_snowflake_sql_api.py @@ -58,6 +58,7 @@ "role": "af_role", }, } + CONN_PARAMS = { "account": "airflow", "application": "AIRFLOW", @@ -71,6 +72,22 @@ "user": "user", "warehouse": "af_wh", } + +CONN_PARAMS_OAUTH = { + "account": "airflow", + "application": "AIRFLOW", + "authenticator": "oauth", + "database": "db", + "client_id": "test_client_id", + "client_secret": "test_client_pw", + "refresh_token": "secrettoken", + "region": "af_region", + "role": "af_role", + "schema": "public", + "session_parameters": None, + "warehouse": "af_wh", +} + HEADERS = { "Content-Type": "application/json", "Authorization": "Bearer newT0k3n", @@ -79,6 +96,15 @@ "X-Snowflake-Authorization-Token-Type": "KEYPAIR_JWT", } +HEADERS_OAUTH = { + "Content-Type": "application/json", + "Authorization": "Bearer newT0k3n", + "Accept": "application/json", + "User-Agent": "snowflakeSQLAPI/1.0", + "X-Snowflake-Authorization-Token-Type": "OAUTH", +} + + GET_RESPONSE = { "resultSetMetaData": { "numRows": 10000, @@ -173,7 +199,7 @@ def test_execute_query( @mock.patch("airflow.providers.snowflake.hooks.snowflake_sql_api.requests") @mock.patch("airflow.providers.snowflake.hooks.snowflake_sql_api.SnowflakeSqlApiHook._get_conn_params") @mock.patch("airflow.providers.snowflake.hooks.snowflake_sql_api.SnowflakeSqlApiHook.get_headers") - def test_execute_query_exception_without_statement_handel( + def test_execute_query_exception_without_statement_handle( self, mock_get_header, mock_conn_param, @@ -250,7 +276,7 @@ def test_get_request_url_header_params(self, mock_get_header, mock_conn_param): @mock.patch("airflow.providers.snowflake.hooks.snowflake_sql_api.SnowflakeSqlApiHook.get_private_key") @mock.patch("airflow.providers.snowflake.hooks.snowflake_sql_api.SnowflakeSqlApiHook._get_conn_params") @mock.patch("airflow.providers.snowflake.utils.sql_api_generate_jwt.JWTGenerator.get_token") - def test_get_headers(self, mock_get_token, mock_conn_param, mock_private_key): + def test_get_headers_should_support_private_key(self, mock_get_token, mock_conn_param, mock_private_key): """Test get_headers method by mocking get_private_key and _get_conn_params method""" mock_get_token.return_value = "newT0k3n" mock_conn_param.return_value = CONN_PARAMS @@ -258,6 +284,38 @@ def test_get_headers(self, mock_get_token, mock_conn_param, mock_private_key): result = hook.get_headers() assert result == HEADERS + @mock.patch("airflow.providers.snowflake.hooks.snowflake_sql_api.SnowflakeSqlApiHook.get_oauth_token") + @mock.patch("airflow.providers.snowflake.hooks.snowflake_sql_api.SnowflakeSqlApiHook._get_conn_params") + def test_get_headers_should_support_oauth(self, mock_conn_param, mock_oauth_token): + """Test get_headers method by mocking get_oauth_token and _get_conn_params method""" + mock_conn_param.return_value = CONN_PARAMS_OAUTH + mock_oauth_token.return_value = "newT0k3n" + hook = SnowflakeSqlApiHook(snowflake_conn_id="mock_conn_id") + result = hook.get_headers() + assert result == HEADERS_OAUTH + + @mock.patch("airflow.providers.snowflake.hooks.snowflake_sql_api.HTTPBasicAuth") + @mock.patch("requests.post") + @mock.patch("airflow.providers.snowflake.hooks.snowflake_sql_api.SnowflakeSqlApiHook._get_conn_params") + def test_get_oauth_token(self, mock_conn_param, requests_post, mock_auth): + """Test get_oauth_token method makes the right http request""" + BASIC_AUTH = {"Authorization": "Basic usernamepassword"} + mock_conn_param.return_value = CONN_PARAMS_OAUTH + requests_post.return_value.status_code = 200 + mock_auth.return_value = BASIC_AUTH + hook = SnowflakeSqlApiHook(snowflake_conn_id="mock_conn_id") + hook.get_oauth_token() + requests_post.assert_called_once_with( + f"https://{CONN_PARAMS_OAUTH['account']}.{CONN_PARAMS_OAUTH['region']}.snowflakecomputing.com/oauth/token-request", + data={ + "grant_type": "refresh_token", + "refresh_token": CONN_PARAMS_OAUTH["refresh_token"], + "redirect_uri": "https://localhost.com", + }, + headers={"Content-Type": "application/x-www-form-urlencoded"}, + auth=BASIC_AUTH, + ) + @pytest.fixture def non_encrypted_temporary_private_key(self, tmp_path: Path) -> Path: """Encrypt the pem file from the path"""