You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "dwreeves (via GitHub)" <gi...@apache.org> on 2026/02/23 19:07:14 UTC
[PR] [Snowflake] [Feat] Allow SnowflakeHook `private_key_content` to use raw key in addition to base64 encoding [airflow]
dwreeves opened a new pull request, #62378:
URL: https://github.com/apache/airflow/pull/62378
Adds support for private keys that are not base64 encoded. This was available prior to the Snowflake provider 6.3.0 release, then it was changed.
There is discussion of this change in #47003 and its corresponding PR. I did not see anyone float the idea of supporting _both_ and a switch was made.
It is argued in the PR that this is not considered a breaking change, rather it is a bug. But in my experience, this is both a breaking change and also not a bug. I use the AWS Secrets Manager backend, not the GUI, and there is no problem on my end storing plain text passwords inside of the AWS Secrets Manager, nor do I have any issues encoding newlines. Migrating to ≥ 6.3.0 would involve significant overhead on my end in not only rotating passwords but changing how they are managed, without a clear migration path which eliminates downtime (precisely because both methods of storing private keys are not supported simultaneously).
This is coming nearly a year later, so perhaps a bit later than it should be if it is to be accepted. But I'm only just doing these migrations right now.
I have an open issue #62359 where this can be discussed further.
* closes: #62359
---
##### Was generative AI tooling used to co-author this PR?
- [ ] Yes (please specify the tool below)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Snowflake] [Feat] Allow SnowflakeHook + SnowflakeSqlApiHook `private_key_content` to use raw key in addition to base64 encoding [airflow]
Posted by "dwreeves (via GitHub)" <gi...@apache.org>.
dwreeves commented on code in PR #62378:
URL: https://github.com/apache/airflow/pull/62378#discussion_r2848722949
##########
providers/snowflake/tests/unit/snowflake/hooks/test_snowflake.py:
##########
@@ -372,6 +372,47 @@ def test_hook_should_support_prepare_basic_conn_params_and_uri(
assert SnowflakeHook(snowflake_conn_id="test_conn").get_uri() == expected_uri
assert SnowflakeHook(snowflake_conn_id="test_conn")._get_conn_params() == expected_conn_params
Review Comment:
Actually, similar to what I did with the other test, I am not going to mock out `b64decode`. My justification is similar to the reasons I gave in the other comment. https://github.com/apache/airflow/pull/62378#discussion_r2847663977
Actually, I have one more justification to add: mocking is annoying because we need to deal with the `MagicMock` object internally, i.e. by hardcoding in the `return_value`. At which point we're trending closer to just testing the test rather than testing the code.
What I did instead was assert that the returned `private_key` is correct:
```python
def test_get_conn_params_should_support_private_auth_in_connection(
self, base64_encoded_encrypted_private_key: str, encrypted_temporary_private_key: Path
):
private_key_content = encrypted_temporary_private_key.read_text()
p_key = serialization.load_pem_private_key(
private_key_content.encode(),
password=_PASSWORD.encode(),
backend=default_backend(),
)
pkb = p_key.private_bytes(
encoding=serialization.Encoding.DER,
format=serialization.PrivateFormat.PKCS8,
encryption_algorithm=serialization.NoEncryption(),
)
connection_kwargs: Any = {
**BASE_CONNECTION_KWARGS,
"password": _PASSWORD,
"extra": {
"database": "db",
"account": "airflow",
"warehouse": "af_wh",
"region": "af_region",
"role": "af_role",
"private_key_content": base64_encoded_encrypted_private_key,
},
}
with (
mock.patch.dict("os.environ", AIRFLOW_CONN_TEST_CONN=Connection(**connection_kwargs).get_uri())
):
conn_params = SnowflakeHook(snowflake_conn_id="test_conn")._get_conn_params()
assert "private_key" in conn_params
assert conn_params["private_key"] == pkb
```
The original code does not do this; all it asserts is `"private_key" in conn_params`.
This is sufficient to ensure that the private key is what we expect it to be: it takes in the b64encoded value, and we ensure it is eventually a valid private key.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Snowflake] [Feat] Allow SnowflakeHook `private_key_content` to use raw key in addition to base64 encoding [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #62378:
URL: https://github.com/apache/airflow/pull/62378#discussion_r2843260344
##########
providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake.py:
##########
@@ -411,7 +411,17 @@ def _get_static_conn_params(self) -> dict[str, str | None]:
raise ValueError("The private_key_file size is too big. Please keep it less than 4 KB.")
private_key_pem = Path(private_key_file_path).read_bytes()
elif private_key_content:
- private_key_pem = base64.b64decode(private_key_content)
+ if any(
+ private_key_content.startswith(header)
+ for header in [
+ "-----BEGIN ENCRYPTED PRIVATE KEY-----\n",
+ "-----BEGIN RSA PRIVATE KEY-----\n",
+ "-----BEGIN PRIVATE KEY-----\n",
Review Comment:
Matching these 3 headers is likely good enough but seems a bit too brittle. And what about cases where the header does not terminate with a new line or has leading/trailing whitespaces. Now, you could argue that it is expected that private keys headers will conform to the set you are validating against but reading RFC 7468, it seems like there is no strict requirements regarding the inclusion of new lines or exclusion of whitespace. Have you considered stripping `private_key_content` of whitespaces and just checking if it starts with `'-----BEGIN'`. Unless there's a need to restrict it to the three PEM formats, I don't think you need to validate to the labels either.
##########
providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake.py:
##########
@@ -411,7 +411,17 @@ def _get_static_conn_params(self) -> dict[str, str | None]:
raise ValueError("The private_key_file size is too big. Please keep it less than 4 KB.")
private_key_pem = Path(private_key_file_path).read_bytes()
elif private_key_content:
- private_key_pem = base64.b64decode(private_key_content)
+ if any(
+ private_key_content.startswith(header)
+ for header in [
+ "-----BEGIN ENCRYPTED PRIVATE KEY-----\n",
+ "-----BEGIN RSA PRIVATE KEY-----\n",
+ "-----BEGIN PRIVATE KEY-----\n",
+ ]
+ ):
+ private_key_pem = private_key_content.encode()
+ else:
+ private_key_pem = base64.b64decode(private_key_content)
Review Comment:
Assuming that the `private_key_content` is in valid base64 format just because it does not start with the headers you enumerated is somewhat of a heuristic. That might likely be true but this part can be made more robust by wrapping line 424 in a try/except and then catching `binascii.Error` and raising an error.
##########
providers/snowflake/tests/unit/snowflake/hooks/test_snowflake.py:
##########
@@ -372,6 +372,47 @@ def test_hook_should_support_prepare_basic_conn_params_and_uri(
assert SnowflakeHook(snowflake_conn_id="test_conn").get_uri() == expected_uri
assert SnowflakeHook(snowflake_conn_id="test_conn")._get_conn_params() == expected_conn_params
+ def test_get_private_key_is_not_base64_encoded_with_headers(
+ self, unencrypted_temporary_private_key: Path, encrypted_temporary_private_key: Path
+ ):
+ """Test get_private_key function skips base64 encoding if private key headers are present."""
+ for pwd, key_file in [
+ (None, unencrypted_temporary_private_key),
+ (_PASSWORD, encrypted_temporary_private_key),
+ ]:
+ private_key_content = key_file.read_text()
+
+ p_key = serialization.load_pem_private_key(
+ private_key_content.encode(),
+ password=(pwd.encode() if pwd is not None else None),
+ backend=default_backend(),
+ )
+
+ pkb = p_key.private_bytes(
+ encoding=serialization.Encoding.DER,
+ format=serialization.PrivateFormat.PKCS8,
+ encryption_algorithm=serialization.NoEncryption(),
+ )
+
+ connection_kwargs: Any = {
+ **BASE_CONNECTION_KWARGS,
+ "password": pwd,
+ "extra": {
+ "database": "db",
+ "account": "airflow",
+ "warehouse": "af_wh",
+ "region": "af_region",
+ "role": "af_role",
+ "private_key_content": private_key_content,
+ },
+ }
+ with mock.patch.dict(
+ "os.environ", AIRFLOW_CONN_TEST_CONN=Connection(**connection_kwargs).get_uri()
+ ):
+ conn_params = SnowflakeHook(snowflake_conn_id="test_conn")._get_conn_params()
+ assert "private_key" in conn_params
+ assert pkb == conn_params["private_key"]
Review Comment:
I would assert that `base64.decode()` is not called. Do not patch it globally. Instead patch it where it is imported in `SnowflakeHook`.
##########
providers/snowflake/tests/unit/snowflake/hooks/test_snowflake.py:
##########
@@ -372,6 +372,47 @@ def test_hook_should_support_prepare_basic_conn_params_and_uri(
assert SnowflakeHook(snowflake_conn_id="test_conn").get_uri() == expected_uri
assert SnowflakeHook(snowflake_conn_id="test_conn")._get_conn_params() == expected_conn_params
Review Comment:
I think you should add another test for the second branch of your implementation i.e. private_key_content does not match the headers in the list. It would be best to assert that `base64.decode()` is not called in that test. Optionally, you could add another test or add a case to the above test for invalid base64 in the same codepath.
##########
providers/snowflake/tests/unit/snowflake/hooks/test_snowflake.py:
##########
@@ -372,6 +372,47 @@ def test_hook_should_support_prepare_basic_conn_params_and_uri(
assert SnowflakeHook(snowflake_conn_id="test_conn").get_uri() == expected_uri
assert SnowflakeHook(snowflake_conn_id="test_conn")._get_conn_params() == expected_conn_params
+ def test_get_private_key_is_not_base64_encoded_with_headers(
+ self, unencrypted_temporary_private_key: Path, encrypted_temporary_private_key: Path
+ ):
+ """Test get_private_key function skips base64 encoding if private key headers are present."""
+ for pwd, key_file in [
+ (None, unencrypted_temporary_private_key),
+ (_PASSWORD, encrypted_temporary_private_key),
+ ]:
Review Comment:
I think it would be better to parametrize this instead of using a loop. This is functional but it makes it difficult to see which branch of the test failed during test runs.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Snowflake] [Feat] Allow SnowflakeHook `private_key_content` to use raw key in addition to base64 encoding [airflow]
Posted by "dwreeves (via GitHub)" <gi...@apache.org>.
dwreeves commented on code in PR #62378:
URL: https://github.com/apache/airflow/pull/62378#discussion_r2847663977
##########
providers/snowflake/tests/unit/snowflake/hooks/test_snowflake.py:
##########
@@ -372,6 +372,47 @@ def test_hook_should_support_prepare_basic_conn_params_and_uri(
assert SnowflakeHook(snowflake_conn_id="test_conn").get_uri() == expected_uri
assert SnowflakeHook(snowflake_conn_id="test_conn")._get_conn_params() == expected_conn_params
+ def test_get_private_key_is_not_base64_encoded_with_headers(
+ self, unencrypted_temporary_private_key: Path, encrypted_temporary_private_key: Path
+ ):
+ """Test get_private_key function skips base64 encoding if private key headers are present."""
+ for pwd, key_file in [
+ (None, unencrypted_temporary_private_key),
+ (_PASSWORD, encrypted_temporary_private_key),
+ ]:
+ private_key_content = key_file.read_text()
+
+ p_key = serialization.load_pem_private_key(
+ private_key_content.encode(),
+ password=(pwd.encode() if pwd is not None else None),
+ backend=default_backend(),
+ )
+
+ pkb = p_key.private_bytes(
+ encoding=serialization.Encoding.DER,
+ format=serialization.PrivateFormat.PKCS8,
+ encryption_algorithm=serialization.NoEncryption(),
+ )
+
+ connection_kwargs: Any = {
+ **BASE_CONNECTION_KWARGS,
+ "password": pwd,
+ "extra": {
+ "database": "db",
+ "account": "airflow",
+ "warehouse": "af_wh",
+ "region": "af_region",
+ "role": "af_role",
+ "private_key_content": private_key_content,
+ },
+ }
+ with mock.patch.dict(
+ "os.environ", AIRFLOW_CONN_TEST_CONN=Connection(**connection_kwargs).get_uri()
+ ):
+ conn_params = SnowflakeHook(snowflake_conn_id="test_conn")._get_conn_params()
+ assert "private_key" in conn_params
+ assert pkb == conn_params["private_key"]
Review Comment:
> Are you worried about another base64.b64decode() call being added in another part of _get_conn_params?
Yes.
I can assert it was either not called or not called with the first arg. If we `try: load_pem_private_key except: b64decode`, then that works.
But-- and sorry for pushing back a bit-- I guess I'm left wondering why it is so important to use `mock` to listen for a call, when it is more explicit just assert the return value is what we expect. Mocking a call and `assert_called` makes a lot of sense for function calls with side-effects or for proper unit test isolation. Here the function can be a black box for all we care; we just want to ensure the private key is read correctly into the params, and I am testing this directly.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Snowflake] [Feat] Allow SnowflakeHook `private_key_content` to use raw key in addition to base64 encoding [airflow]
Posted by "dwreeves (via GitHub)" <gi...@apache.org>.
dwreeves commented on PR #62378:
URL: https://github.com/apache/airflow/pull/62378#issuecomment-3952696083
Thank you @SameerMesiah97 for @bugraoz93 for your feedback and helping me get this PR in a much better state. I'm continuing to work on this.
I also realize that the `SnowflakeSqlApiHook` has a separate private key flow, so I'll add support to that. I'm also wondering if those should be consolidated to be more DRY, and whether that would be out of scope. 🤔
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Snowflake] [Feat] Allow SnowflakeHook `private_key_content` to use raw key in addition to base64 encoding [airflow]
Posted by "dwreeves (via GitHub)" <gi...@apache.org>.
dwreeves commented on code in PR #62378:
URL: https://github.com/apache/airflow/pull/62378#discussion_r2847667899
##########
providers/snowflake/tests/unit/snowflake/hooks/test_snowflake.py:
##########
@@ -372,6 +372,47 @@ def test_hook_should_support_prepare_basic_conn_params_and_uri(
assert SnowflakeHook(snowflake_conn_id="test_conn").get_uri() == expected_uri
assert SnowflakeHook(snowflake_conn_id="test_conn")._get_conn_params() == expected_conn_params
+ def test_get_private_key_is_not_base64_encoded_with_headers(
+ self, unencrypted_temporary_private_key: Path, encrypted_temporary_private_key: Path
+ ):
+ """Test get_private_key function skips base64 encoding if private key headers are present."""
+ for pwd, key_file in [
+ (None, unencrypted_temporary_private_key),
+ (_PASSWORD, encrypted_temporary_private_key),
+ ]:
Review Comment:
Makes sense, I'll do two separate tests.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Snowflake] [Feat] Allow SnowflakeHook + SnowflakeSqlApiHook `private_key_content` to use raw key in addition to base64 encoding [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #62378:
URL: https://github.com/apache/airflow/pull/62378#issuecomment-4026614727
> But I'm not sure if that's necessary. get_private_key() is essentially a private method (not marked as such, but end-users would not realistically ever call it) and it is only called once, so I modified the code at that point to be self.private_key = self.get_private_key() instead of setting self.private_key as a side-effect.
Not necessary. This looks good :)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Snowflake] [Feat] Allow SnowflakeHook `private_key_content` to use raw key in addition to base64 encoding [airflow]
Posted by "bugraoz93 (via GitHub)" <gi...@apache.org>.
bugraoz93 commented on code in PR #62378:
URL: https://github.com/apache/airflow/pull/62378#discussion_r2843233868
##########
providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake.py:
##########
@@ -411,7 +411,17 @@ def _get_static_conn_params(self) -> dict[str, str | None]:
raise ValueError("The private_key_file size is too big. Please keep it less than 4 KB.")
private_key_pem = Path(private_key_file_path).read_bytes()
elif private_key_content:
- private_key_pem = base64.b64decode(private_key_content)
+ if any(
Review Comment:
How about something like this rather than string-searching static constants? This way you can also verify the content is loadable using `cryptography` which is coming from Snowflake connector and used within this file :)
https://cryptography.io/en/latest/hazmat/primitives/asymmetric/serialization/#cryptography.hazmat.primitives.serialization.load_pem_private_key
``` Python
from cryptography.hazmat.primitives.serialization import load_pem_private_key
def __is_loadable_pem(self, content: str):
try:
load_pem_private_key(content, password=password)
return True
except (ValueError, TypeError):
return False
```
``` Python
if self.__is_loadable_pem(private_key_content):
private_key_pem = private_key_content.encode()
else:
private_key_pem = base64.b64decode(private_key_content)
```
Maybe rather than return false, we can also handle exceptions from exceptions in base64 decoding and exceptions in pem file validation, then convert accordingly. It could be too much for the PR. We can follow up as an improvement too
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Snowflake] [Feat] Allow SnowflakeHook `private_key_content` to use raw key in addition to base64 encoding [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #62378:
URL: https://github.com/apache/airflow/pull/62378#discussion_r2847725086
##########
providers/snowflake/tests/unit/snowflake/hooks/test_snowflake.py:
##########
@@ -372,6 +372,47 @@ def test_hook_should_support_prepare_basic_conn_params_and_uri(
assert SnowflakeHook(snowflake_conn_id="test_conn").get_uri() == expected_uri
assert SnowflakeHook(snowflake_conn_id="test_conn")._get_conn_params() == expected_conn_params
+ def test_get_private_key_is_not_base64_encoded_with_headers(
+ self, unencrypted_temporary_private_key: Path, encrypted_temporary_private_key: Path
+ ):
+ """Test get_private_key function skips base64 encoding if private key headers are present."""
+ for pwd, key_file in [
+ (None, unencrypted_temporary_private_key),
+ (_PASSWORD, encrypted_temporary_private_key),
+ ]:
+ private_key_content = key_file.read_text()
+
+ p_key = serialization.load_pem_private_key(
+ private_key_content.encode(),
+ password=(pwd.encode() if pwd is not None else None),
+ backend=default_backend(),
+ )
+
+ pkb = p_key.private_bytes(
+ encoding=serialization.Encoding.DER,
+ format=serialization.PrivateFormat.PKCS8,
+ encryption_algorithm=serialization.NoEncryption(),
+ )
+
+ connection_kwargs: Any = {
+ **BASE_CONNECTION_KWARGS,
+ "password": pwd,
+ "extra": {
+ "database": "db",
+ "account": "airflow",
+ "warehouse": "af_wh",
+ "region": "af_region",
+ "role": "af_role",
+ "private_key_content": private_key_content,
+ },
+ }
+ with mock.patch.dict(
+ "os.environ", AIRFLOW_CONN_TEST_CONN=Connection(**connection_kwargs).get_uri()
+ ):
+ conn_params = SnowflakeHook(snowflake_conn_id="test_conn")._get_conn_params()
+ assert "private_key" in conn_params
+ assert pkb == conn_params["private_key"]
Review Comment:
> > Are you worried about another base64.b64decode() call being added in another part of _get_conn_params?
>
> Yes.
>
> I can assert it was either not called or not called with the first arg. If we `try: load_pem_private_key except: b64decode`, then that works.
>
> But-- and sorry for pushing back a bit-- I guess I'm left wondering why it is so important to use `mock` to listen for a call, when it is more explicit just assert the return value is what we expect. Mocking a call and `assert_called` makes a lot of sense for function calls with side-effects or for proper unit test isolation. Here the function can be a black box for all we care; we just want to ensure the private key is read correctly into the params, and I am testing this directly.
Okay. I see your point. I agree it is not critical to the test but it would make the implementation more robust. The PR should not be blocked because of this. It's up to you accept the suggestion. As long as the return value is being checked, we should be fine.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Snowflake] [Feat] Allow SnowflakeHook `private_key_content` to use raw key in addition to base64 encoding [airflow]
Posted by "dwreeves (via GitHub)" <gi...@apache.org>.
dwreeves commented on code in PR #62378:
URL: https://github.com/apache/airflow/pull/62378#discussion_r2847390823
##########
providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake.py:
##########
@@ -411,7 +411,17 @@ def _get_static_conn_params(self) -> dict[str, str | None]:
raise ValueError("The private_key_file size is too big. Please keep it less than 4 KB.")
private_key_pem = Path(private_key_file_path).read_bytes()
elif private_key_content:
- private_key_pem = base64.b64decode(private_key_content)
+ if any(
+ private_key_content.startswith(header)
+ for header in [
+ "-----BEGIN ENCRYPTED PRIVATE KEY-----\n",
+ "-----BEGIN RSA PRIVATE KEY-----\n",
+ "-----BEGIN PRIVATE KEY-----\n",
Review Comment:
I'm going to go with what @bugraoz93 suggested and do a try/except EAFP approach instead of this more LBYL approach, if that's OK.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Snowflake] [Feat] Allow SnowflakeHook `private_key_content` to use raw key in addition to base64 encoding [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #62378:
URL: https://github.com/apache/airflow/pull/62378#discussion_r2847603850
##########
providers/snowflake/tests/unit/snowflake/hooks/test_snowflake.py:
##########
@@ -372,6 +372,47 @@ def test_hook_should_support_prepare_basic_conn_params_and_uri(
assert SnowflakeHook(snowflake_conn_id="test_conn").get_uri() == expected_uri
assert SnowflakeHook(snowflake_conn_id="test_conn")._get_conn_params() == expected_conn_params
+ def test_get_private_key_is_not_base64_encoded_with_headers(
+ self, unencrypted_temporary_private_key: Path, encrypted_temporary_private_key: Path
+ ):
+ """Test get_private_key function skips base64 encoding if private key headers are present."""
+ for pwd, key_file in [
+ (None, unencrypted_temporary_private_key),
+ (_PASSWORD, encrypted_temporary_private_key),
+ ]:
+ private_key_content = key_file.read_text()
+
+ p_key = serialization.load_pem_private_key(
+ private_key_content.encode(),
+ password=(pwd.encode() if pwd is not None else None),
+ backend=default_backend(),
+ )
+
+ pkb = p_key.private_bytes(
+ encoding=serialization.Encoding.DER,
+ format=serialization.PrivateFormat.PKCS8,
+ encryption_algorithm=serialization.NoEncryption(),
+ )
+
+ connection_kwargs: Any = {
+ **BASE_CONNECTION_KWARGS,
+ "password": pwd,
+ "extra": {
+ "database": "db",
+ "account": "airflow",
+ "warehouse": "af_wh",
+ "region": "af_region",
+ "role": "af_role",
+ "private_key_content": private_key_content,
+ },
+ }
+ with mock.patch.dict(
+ "os.environ", AIRFLOW_CONN_TEST_CONN=Connection(**connection_kwargs).get_uri()
+ ):
+ conn_params = SnowflakeHook(snowflake_conn_id="test_conn")._get_conn_params()
+ assert "private_key" in conn_params
+ assert pkb == conn_params["private_key"]
Review Comment:
> Maybe this is a little overkill, but I made a deliberate decision to avoid mocking `base64.decode()` because I am unclear if, in the future, `base64` may be called inside the function, but outside of the private key flow. So I felt there was a very small but nonzero chance that asserting `base64.decode()` could, sometime in the future, incorrectly fail the test, and that this approach was more explicit.
>
> WDYT?
That is understandable. Are you worried about another `base64.b64decode()` call being added in another part of `_get_conn_params`? Maybe you could instead assert that `base64.b64decode()` was not called with `private_key_content`? There is no in-built functionality to do this but perhaps you could assert that call for `private_key_content`is not in `mock.call_args_list` after mocking `base64.b64decode()`?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Snowflake] [Feat] Allow SnowflakeHook `private_key_content` to use raw key in addition to base64 encoding [airflow]
Posted by "dwreeves (via GitHub)" <gi...@apache.org>.
dwreeves commented on code in PR #62378:
URL: https://github.com/apache/airflow/pull/62378#discussion_r2843322351
##########
providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake.py:
##########
@@ -411,7 +411,17 @@ def _get_static_conn_params(self) -> dict[str, str | None]:
raise ValueError("The private_key_file size is too big. Please keep it less than 4 KB.")
private_key_pem = Path(private_key_file_path).read_bytes()
elif private_key_content:
- private_key_pem = base64.b64decode(private_key_content)
+ if any(
Review Comment:
Sure, that works. I wasn't sure if a if/else or try/except was more appropriate. I can implement the latter.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Snowflake] [Feat] Allow SnowflakeHook `private_key_content` to use raw key in addition to base64 encoding [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #62378:
URL: https://github.com/apache/airflow/pull/62378#discussion_r2843353553
##########
providers/snowflake/tests/unit/snowflake/hooks/test_snowflake.py:
##########
@@ -372,6 +372,47 @@ def test_hook_should_support_prepare_basic_conn_params_and_uri(
assert SnowflakeHook(snowflake_conn_id="test_conn").get_uri() == expected_uri
assert SnowflakeHook(snowflake_conn_id="test_conn")._get_conn_params() == expected_conn_params
Review Comment:
I think you should add another test for the second branch of your implementation i.e. private_key_content does not match the headers in the list. It would be best to assert that `base64.decode()` is called in that test. Optionally, you could add another test or add a case to the above test for invalid base64 in the same codepath.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Snowflake] [Feat] Allow SnowflakeHook `private_key_content` to use raw key in addition to base64 encoding [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on PR #62378:
URL: https://github.com/apache/airflow/pull/62378#issuecomment-3952986463
> Thank you @SameerMesiah97 for @bugraoz93 for your feedback and helping me get this PR in a much better state. I'm continuing to work on this.
>
> I also realize that the `SnowflakeSqlApiHook` has a separate private key flow, so I'll add support to that. I'm also wondering if those should be consolidated to be more DRY, and whether that would be out of scope. 🤔
If you wish to do that, I would broaden the scope for this issue to include `SnowflakeSqlApiHook` as well (update the description). And I think we should see how the private key flow looks in `SnowflakeSqlApiHook` before considering consolidating in say the `utils` module for the provider. Note that `SnowflakeSqlApiHook` inherits from `SnowflakeHook`and whenever it needs to use functionality from there, it usually either calls the method via self (see `_get_conn_params`) or uses `super() `on the relevant method alongside any injected code (see `get_oauth_token`). I think should wait to see how your implementation for `SnowflakeSqlApiHook` looks like to see how it should be handled.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Snowflake] [Feat] Allow SnowflakeHook `private_key_content` to use raw key in addition to base64 encoding [airflow]
Posted by "dwreeves (via GitHub)" <gi...@apache.org>.
dwreeves commented on code in PR #62378:
URL: https://github.com/apache/airflow/pull/62378#discussion_r2847394176
##########
providers/snowflake/tests/unit/snowflake/hooks/test_snowflake.py:
##########
@@ -372,6 +372,47 @@ def test_hook_should_support_prepare_basic_conn_params_and_uri(
assert SnowflakeHook(snowflake_conn_id="test_conn").get_uri() == expected_uri
assert SnowflakeHook(snowflake_conn_id="test_conn")._get_conn_params() == expected_conn_params
+ def test_get_private_key_is_not_base64_encoded_with_headers(
+ self, unencrypted_temporary_private_key: Path, encrypted_temporary_private_key: Path
+ ):
+ """Test get_private_key function skips base64 encoding if private key headers are present."""
+ for pwd, key_file in [
+ (None, unencrypted_temporary_private_key),
+ (_PASSWORD, encrypted_temporary_private_key),
+ ]:
Review Comment:
Good point, I'll use pytest mark parameterize.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Snowflake] [Feat] Allow SnowflakeHook + SnowflakeSqlApiHook `private_key_content` to use raw key in addition to base64 encoding [airflow]
Posted by "dwreeves (via GitHub)" <gi...@apache.org>.
dwreeves commented on PR #62378:
URL: https://github.com/apache/airflow/pull/62378#issuecomment-3953895699
OK, assuming CI passes, I think we're at a good spot.
- Added support for plain text secret key, which makes the upgrade path from <6.3.0 much easier, while still supporting ≥6.3.0 behavior of b64encoded private keys.
- Updated docs.
- Excluding tests, net negative lines of code via removing repetitive code which handles the private key. Neat!
---
**The only unknown**, and it would be helpful for a maintainer to chime in 😄, is the question of whether there are reasons _other than the UI_ that private keys in plain text are no longer supported. Based on looking through discussions, I did not see any reasons other than the UI. I'd be interested in hearing if there are other reasons. If not, I would argue that this is an appropriate feature addition.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Snowflake] [Feat] Allow SnowflakeHook + SnowflakeSqlApiHook `private_key_content` to use raw key in addition to base64 encoding [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk merged PR #62378:
URL: https://github.com/apache/airflow/pull/62378
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Snowflake] [Feat] Allow SnowflakeHook `private_key_content` to use raw key in addition to base64 encoding [airflow]
Posted by "dwreeves (via GitHub)" <gi...@apache.org>.
dwreeves commented on PR #62378:
URL: https://github.com/apache/airflow/pull/62378#issuecomment-3953247376
> I think should wait to see how your implementation for SnowflakeSqlApiHook looks like to see how it should be handled.
I just pushed the change, so you can see.
I moved `get_private_key()` to `SnowflakeHook`, and gave it a return value instead of setting `self.private_key`.
If you want true, pure backwards compat, then we can have `SnowflakeSqlApiHook` have:
```python
def get_private_key(self):
self.private_key = super().get_private_key()
return self.private_key
```
But I'm not sure if that's necessary. `get_private_key()` is essentially a private method (not marked as such, but end-users would not realistically ever call it) and it is only called once, so I modified the code at that point to be `self.private_key = self.get_private_key()` instead of setting `self.private_key` as a side-effect.
(If you treat `get_private_key()` literally as a public method, this is a minorly breaking change, but I really have a hard time seeing it as a true public method, and often people writing Python code aren't so strict about such decisions so I'm OK interpreting it as private.)
I did not bother to put inside of `utils`. I'm not against it, I have no opinion and I don't want to go crazy touching everything for code I don't maintain. As far as I see it, this code is better and DRYer than I left it, and I see what's here already as an improvement. Let me know if you think it should go in `utils` and what that would look like. I do see you touch this provider package quite a bit so you can say better than me.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Snowflake] [Feat] Allow SnowflakeHook `private_key_content` to use raw key in addition to base64 encoding [airflow]
Posted by "dwreeves (via GitHub)" <gi...@apache.org>.
dwreeves commented on code in PR #62378:
URL: https://github.com/apache/airflow/pull/62378#discussion_r2847409866
##########
providers/snowflake/tests/unit/snowflake/hooks/test_snowflake.py:
##########
@@ -372,6 +372,47 @@ def test_hook_should_support_prepare_basic_conn_params_and_uri(
assert SnowflakeHook(snowflake_conn_id="test_conn").get_uri() == expected_uri
assert SnowflakeHook(snowflake_conn_id="test_conn")._get_conn_params() == expected_conn_params
+ def test_get_private_key_is_not_base64_encoded_with_headers(
+ self, unencrypted_temporary_private_key: Path, encrypted_temporary_private_key: Path
+ ):
+ """Test get_private_key function skips base64 encoding if private key headers are present."""
+ for pwd, key_file in [
+ (None, unencrypted_temporary_private_key),
+ (_PASSWORD, encrypted_temporary_private_key),
+ ]:
+ private_key_content = key_file.read_text()
+
+ p_key = serialization.load_pem_private_key(
+ private_key_content.encode(),
+ password=(pwd.encode() if pwd is not None else None),
+ backend=default_backend(),
+ )
+
+ pkb = p_key.private_bytes(
+ encoding=serialization.Encoding.DER,
+ format=serialization.PrivateFormat.PKCS8,
+ encryption_algorithm=serialization.NoEncryption(),
+ )
+
+ connection_kwargs: Any = {
+ **BASE_CONNECTION_KWARGS,
+ "password": pwd,
+ "extra": {
+ "database": "db",
+ "account": "airflow",
+ "warehouse": "af_wh",
+ "region": "af_region",
+ "role": "af_role",
+ "private_key_content": private_key_content,
+ },
+ }
+ with mock.patch.dict(
+ "os.environ", AIRFLOW_CONN_TEST_CONN=Connection(**connection_kwargs).get_uri()
+ ):
+ conn_params = SnowflakeHook(snowflake_conn_id="test_conn")._get_conn_params()
+ assert "private_key" in conn_params
+ assert pkb == conn_params["private_key"]
Review Comment:
Maybe this is a little overkill, but I made a deliberate decision to avoid mocking `base64.decode()` because I am unclear if, in the future, `base64` may be called inside the function, but outside of the private key flow. So I felt there was a very small but nonzero chance that asserting `base64.decode()` could, sometime in the future, incorrectly fail the test, and that this approach was more explicit.
WDYT?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Snowflake] [Feat] Allow SnowflakeHook `private_key_content` to use raw key in addition to base64 encoding [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #62378:
URL: https://github.com/apache/airflow/pull/62378#discussion_r2847480637
##########
providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake.py:
##########
@@ -411,7 +411,17 @@ def _get_static_conn_params(self) -> dict[str, str | None]:
raise ValueError("The private_key_file size is too big. Please keep it less than 4 KB.")
private_key_pem = Path(private_key_file_path).read_bytes()
elif private_key_content:
- private_key_pem = base64.b64decode(private_key_content)
+ if any(
+ private_key_content.startswith(header)
+ for header in [
+ "-----BEGIN ENCRYPTED PRIVATE KEY-----\n",
+ "-----BEGIN RSA PRIVATE KEY-----\n",
+ "-----BEGIN PRIVATE KEY-----\n",
Review Comment:
Sure.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Snowflake] [Feat] Allow SnowflakeHook `private_key_content` to use raw key in addition to base64 encoding [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #62378:
URL: https://github.com/apache/airflow/pull/62378#discussion_r2847611450
##########
providers/snowflake/tests/unit/snowflake/hooks/test_snowflake.py:
##########
@@ -372,6 +372,47 @@ def test_hook_should_support_prepare_basic_conn_params_and_uri(
assert SnowflakeHook(snowflake_conn_id="test_conn").get_uri() == expected_uri
assert SnowflakeHook(snowflake_conn_id="test_conn")._get_conn_params() == expected_conn_params
+ def test_get_private_key_is_not_base64_encoded_with_headers(
+ self, unencrypted_temporary_private_key: Path, encrypted_temporary_private_key: Path
+ ):
+ """Test get_private_key function skips base64 encoding if private key headers are present."""
+ for pwd, key_file in [
+ (None, unencrypted_temporary_private_key),
+ (_PASSWORD, encrypted_temporary_private_key),
+ ]:
Review Comment:
Okay. Do 2 separate tests then.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Snowflake] [Feat] Allow SnowflakeHook `private_key_content` to use raw key in addition to base64 encoding [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #62378:
URL: https://github.com/apache/airflow/pull/62378#discussion_r2843353553
##########
providers/snowflake/tests/unit/snowflake/hooks/test_snowflake.py:
##########
@@ -372,6 +372,47 @@ def test_hook_should_support_prepare_basic_conn_params_and_uri(
assert SnowflakeHook(snowflake_conn_id="test_conn").get_uri() == expected_uri
assert SnowflakeHook(snowflake_conn_id="test_conn")._get_conn_params() == expected_conn_params
Review Comment:
I think you should add another test for the second branch of your implementation i.e. private_key_content does not match the headers in the list. It would be best to assert that `base64.b64decode()` is called in that test. Optionally, you could add another test or add a case to the above test for invalid base64 in the same codepath.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Snowflake] [Feat] Allow SnowflakeHook + SnowflakeSqlApiHook `private_key_content` to use raw key in addition to base64 encoding [airflow]
Posted by "dwreeves (via GitHub)" <gi...@apache.org>.
dwreeves commented on code in PR #62378:
URL: https://github.com/apache/airflow/pull/62378#discussion_r2848610533
##########
providers/snowflake/tests/unit/snowflake/hooks/test_snowflake.py:
##########
@@ -372,6 +372,47 @@ def test_hook_should_support_prepare_basic_conn_params_and_uri(
assert SnowflakeHook(snowflake_conn_id="test_conn").get_uri() == expected_uri
assert SnowflakeHook(snowflake_conn_id="test_conn")._get_conn_params() == expected_conn_params
Review Comment:
I feel like this is partly though not fully covered already by existing tests like `test_get_conn_params_should_support_private_auth_in_connection`.
What I think may make sense is to just assert that `b64decode` is called in those existing tests. But, additional requests aren't required, per se, just additional asserts in existing ones.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Snowflake] [Feat] Allow SnowflakeHook + SnowflakeSqlApiHook `private_key_content` to use raw key in addition to base64 encoding [airflow]
Posted by "dwreeves (via GitHub)" <gi...@apache.org>.
dwreeves commented on code in PR #62378:
URL: https://github.com/apache/airflow/pull/62378#discussion_r2848722949
##########
providers/snowflake/tests/unit/snowflake/hooks/test_snowflake.py:
##########
@@ -372,6 +372,47 @@ def test_hook_should_support_prepare_basic_conn_params_and_uri(
assert SnowflakeHook(snowflake_conn_id="test_conn").get_uri() == expected_uri
assert SnowflakeHook(snowflake_conn_id="test_conn")._get_conn_params() == expected_conn_params
Review Comment:
Actually, similar to what I did with the other test, I am not going to mock out `b64decode`. My justification is similar to the reasons I gave in the other comment. https://github.com/apache/airflow/pull/62378#discussion_r2847663977
Actually, I have one more justification to add: mocking is annoying because we need to deal with the `MagicMock` object internally, i.e. by hardcoding in the `return_value`. At which point we're trending closer to just testing the test rather than testing the code.
What I did instead was assert that the returned `private_key` is correct:
```python
def test_get_conn_params_should_support_private_auth_in_connection(
self, base64_encoded_encrypted_private_key: str, encrypted_temporary_private_key: Path
):
private_key_content = encrypted_temporary_private_key.read_text()
p_key = serialization.load_pem_private_key(
private_key_content.encode(),
password=_PASSWORD.encode(),
backend=default_backend(),
)
pkb = p_key.private_bytes(
encoding=serialization.Encoding.DER,
format=serialization.PrivateFormat.PKCS8,
encryption_algorithm=serialization.NoEncryption(),
)
connection_kwargs: Any = {
**BASE_CONNECTION_KWARGS,
"password": _PASSWORD,
"extra": {
"database": "db",
"account": "airflow",
"warehouse": "af_wh",
"region": "af_region",
"role": "af_role",
"private_key_content": base64_encoded_encrypted_private_key,
},
}
with (
mock.patch.dict("os.environ", AIRFLOW_CONN_TEST_CONN=Connection(**connection_kwargs).get_uri())
):
conn_params = SnowflakeHook(snowflake_conn_id="test_conn")._get_conn_params()
assert "private_key" in conn_params
assert conn_params["private_key"] == pkb
```
This is sufficient to ensure that the private key is what we expect it to be: it takes in the b64encoded value, and we ensure it is eventually a valid private key.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Snowflake] [Feat] Allow SnowflakeHook `private_key_content` to use raw key in addition to base64 encoding [airflow]
Posted by "dwreeves (via GitHub)" <gi...@apache.org>.
dwreeves commented on code in PR #62378:
URL: https://github.com/apache/airflow/pull/62378#discussion_r2847512608
##########
providers/snowflake/tests/unit/snowflake/hooks/test_snowflake.py:
##########
@@ -372,6 +372,47 @@ def test_hook_should_support_prepare_basic_conn_params_and_uri(
assert SnowflakeHook(snowflake_conn_id="test_conn").get_uri() == expected_uri
assert SnowflakeHook(snowflake_conn_id="test_conn")._get_conn_params() == expected_conn_params
+ def test_get_private_key_is_not_base64_encoded_with_headers(
+ self, unencrypted_temporary_private_key: Path, encrypted_temporary_private_key: Path
+ ):
+ """Test get_private_key function skips base64 encoding if private key headers are present."""
+ for pwd, key_file in [
+ (None, unencrypted_temporary_private_key),
+ (_PASSWORD, encrypted_temporary_private_key),
+ ]:
Review Comment:
So after saying I would use `pytest.mark.parametrize`, I realize a simple problem here: passing in other fixtures to `pytest.mark.parametrize` is somewhat nontrivial. All the "nice" solutions I can think of require more rewrite than this feature change is worth. Let me know what you think is the best path forward. I could just do two separate tests.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Snowflake] [Feat] Allow SnowflakeHook + SnowflakeSqlApiHook `private_key_content` to use raw key in addition to base64 encoding [airflow]
Posted by "dwreeves (via GitHub)" <gi...@apache.org>.
dwreeves commented on code in PR #62378:
URL: https://github.com/apache/airflow/pull/62378#discussion_r2848244149
##########
providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake.py:
##########
@@ -411,7 +411,17 @@ def _get_static_conn_params(self) -> dict[str, str | None]:
raise ValueError("The private_key_file size is too big. Please keep it less than 4 KB.")
private_key_pem = Path(private_key_file_path).read_bytes()
elif private_key_content:
- private_key_pem = base64.b64decode(private_key_content)
+ if any(
+ private_key_content.startswith(header)
+ for header in [
+ "-----BEGIN ENCRYPTED PRIVATE KEY-----\n",
+ "-----BEGIN RSA PRIVATE KEY-----\n",
+ "-----BEGIN PRIVATE KEY-----\n",
Review Comment:
Done.
##########
providers/snowflake/tests/unit/snowflake/hooks/test_snowflake.py:
##########
@@ -372,6 +372,47 @@ def test_hook_should_support_prepare_basic_conn_params_and_uri(
assert SnowflakeHook(snowflake_conn_id="test_conn").get_uri() == expected_uri
assert SnowflakeHook(snowflake_conn_id="test_conn")._get_conn_params() == expected_conn_params
+ def test_get_private_key_is_not_base64_encoded_with_headers(
+ self, unencrypted_temporary_private_key: Path, encrypted_temporary_private_key: Path
+ ):
+ """Test get_private_key function skips base64 encoding if private key headers are present."""
+ for pwd, key_file in [
+ (None, unencrypted_temporary_private_key),
+ (_PASSWORD, encrypted_temporary_private_key),
+ ]:
Review Comment:
Done.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Snowflake] [Feat] Allow SnowflakeHook + SnowflakeSqlApiHook `private_key_content` to use raw key in addition to base64 encoding [airflow]
Posted by "dwreeves (via GitHub)" <gi...@apache.org>.
dwreeves commented on code in PR #62378:
URL: https://github.com/apache/airflow/pull/62378#discussion_r2848241791
##########
providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake.py:
##########
@@ -411,7 +411,17 @@ def _get_static_conn_params(self) -> dict[str, str | None]:
raise ValueError("The private_key_file size is too big. Please keep it less than 4 KB.")
private_key_pem = Path(private_key_file_path).read_bytes()
elif private_key_content:
- private_key_pem = base64.b64decode(private_key_content)
+ if any(
Review Comment:
Updated.
I didn't use the private method, instead I just did a try-except directly.
The upside of direct try-except is this prevents repeated work-- extremely minor repeated work, but still-- and keeps the control flow self-contained.
The downside is it makes the traceback slightly more annoying to parse through, and it's probably not terribly uncommon for users to hit a b64encoding error here so this longer error traceback will be user facing (albeit, the traceback is arguably more informative, so not necessarily even a downside!).
Let me know if you think it should be its own method. If you're OK with it not being one, then this is complete.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org