Moved sql_api_generate_jwt out of hooks folder (#31586)

* Moved sql_api_generate_jwt out of hooks folder
apache · May 29, 2023 · c55f9ea · c55f9ea
1 parent 55d3ecc
commit c55f9ea
Show file tree

Hide file tree

Showing 5 changed files with 3 additions and 4 deletions.
diff --git a/airflow/providers/snowflake/hooks/snowflake_sql_api.py b/airflow/providers/snowflake/hooks/snowflake_sql_api.py
@@ -27,7 +27,7 @@
 
 from airflow import AirflowException
 from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook
-from airflow.providers.snowflake.hooks.sql_api_generate_jwt import JWTGenerator
+from airflow.providers.snowflake.utils.sql_api_generate_jwt import JWTGenerator
 
 
 class SnowflakeSqlApiHook(SnowflakeHook):

diff --git a/airflow/providers/snowflake/provider.yaml b/airflow/providers/snowflake/provider.yaml
@@ -76,7 +76,6 @@ hooks:
     python-modules:
       - airflow.providers.snowflake.hooks.snowflake
       - airflow.providers.snowflake.hooks.snowflake_sql_api
-      - airflow.providers.snowflake.hooks.sql_api_generate_jwt
 
 transfers:
   - source-integration-name: Amazon Simple Storage Service (S3)

diff --git a/...s/snowflake/hooks/sql_api_generate_jwt.py → ...s/snowflake/utils/sql_api_generate_jwt.py b/...s/snowflake/hooks/sql_api_generate_jwt.py → ...s/snowflake/utils/sql_api_generate_jwt.py
diff --git a/tests/providers/snowflake/hooks/test_snowflake_sql_api.py b/tests/providers/snowflake/hooks/test_snowflake_sql_api.py
@@ -254,7 +254,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.hooks.sql_api_generate_jwt.JWTGenerator.get_token")
+    @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):
         """Test get_headers method by mocking get_private_key and _get_conn_params method"""
         mock_get_token.return_value = "newT0k3n"

diff --git a/...wflake/hooks/test_sql_api_generate_jwt.py → ...wflake/utils/test_sql_api_generate_jwt.py b/...wflake/hooks/test_sql_api_generate_jwt.py → ...wflake/utils/test_sql_api_generate_jwt.py
@@ -21,7 +21,7 @@
 from cryptography.hazmat.primitives import serialization as crypto_serialization
 from cryptography.hazmat.primitives.asymmetric import rsa
 
-from airflow.providers.snowflake.hooks.sql_api_generate_jwt import JWTGenerator
+from airflow.providers.snowflake.utils.sql_api_generate_jwt import JWTGenerator
 
 _PASSWORD = "snowflake42"