You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "utkarsharma2 (via GitHub)" <gi...@apache.org> on 2023/05/09 15:04:18 UTC
[GitHub] [airflow] utkarsharma2 opened a new pull request, #30698: Add SnowflakeSqlApiOperator operator
utkarsharma2 opened a new pull request, #30698:
URL: https://github.com/apache/airflow/pull/30698
Implemented Snowflake SQL API Operator to support multiple SQL statements, the Snowflake SQL API allows submitting multiple SQL statements in a single request, which currently is not possible with `SnowflakeOperator` as it submits multiple statements sequentially.
This Operator currently uses key pair authentication, so you need to provide private key raw content or private key file path in the snowflake connection along with other details.
Snowflake SQL API key pair - Authentication https://docs.snowflake.com/en/developer-guide/sql-api/authenticating.html#label-sql-api-authenticating-key-pair
--
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
[GitHub] [airflow] ephraimbuddy merged pull request #30698: Add SnowflakeSqlApiOperator operator
Posted by "ephraimbuddy (via GitHub)" <gi...@apache.org>.
ephraimbuddy merged PR #30698:
URL: https://github.com/apache/airflow/pull/30698
--
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
[GitHub] [airflow] eladkal commented on a diff in pull request #30698: Add SnowflakeSqlApiOperator operator
Posted by "eladkal (via GitHub)" <gi...@apache.org>.
eladkal commented on code in PR #30698:
URL: https://github.com/apache/airflow/pull/30698#discussion_r1212154907
##########
airflow/providers/snowflake/operators/snowflake.py:
##########
@@ -359,3 +368,159 @@ def __init__(
self.authenticator = authenticator
self.session_parameters = session_parameters
self.query_ids: list[str] = []
+
+
+class SnowflakeSqlApiOperator(SnowflakeOperator):
Review Comment:
This is not good.
`SnowflakeOperator` is deprecated. We should not extend it's functionality.
This class needs to be refactored around `SQLExecuteQueryOperator` or `BaseOperator` ?
--
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
[GitHub] [airflow] eladkal commented on a diff in pull request #30698: Add SnowflakeSqlApiOperator operator
Posted by "eladkal (via GitHub)" <gi...@apache.org>.
eladkal commented on code in PR #30698:
URL: https://github.com/apache/airflow/pull/30698#discussion_r1212154907
##########
airflow/providers/snowflake/operators/snowflake.py:
##########
@@ -359,3 +368,159 @@ def __init__(
self.authenticator = authenticator
self.session_parameters = session_parameters
self.query_ids: list[str] = []
+
+
+class SnowflakeSqlApiOperator(SnowflakeOperator):
Review Comment:
This is not good.
`SnowflakeOperator` is deprecated. We should not extend it's functionality.
This class needs to be refactored around `SQLExecuteQueryOperator`
--
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
[GitHub] [airflow] utkarsharma2 commented on a diff in pull request #30698: Add SnowflakeSqlApiOperator operator
Posted by "utkarsharma2 (via GitHub)" <gi...@apache.org>.
utkarsharma2 commented on code in PR #30698:
URL: https://github.com/apache/airflow/pull/30698#discussion_r1189326460
##########
airflow/providers/snowflake/hooks/snowflake_sql_api.py:
##########
@@ -0,0 +1,258 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import uuid
+from datetime import timedelta
+from pathlib import Path
+from typing import Any
+
+import requests
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives import serialization
+
+from airflow import AirflowException
+from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook
+from airflow.providers.snowflake.hooks.sql_api_generate_jwt import JWTGenerator
+
+
+class SnowflakeSqlApiHook(SnowflakeHook):
+ """
+ A client to interact with Snowflake using SQL API and allows submitting
+ multiple SQL statements in a single request. In combination with aiohttp, make post request to submit SQL
+ statements for execution, 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.
+
+ :param snowflake_conn_id: Reference to
+ :ref:`Snowflake connection id<howto/connection:snowflake>`
+ :param account: snowflake account name
+ :param authenticator: authenticator for Snowflake.
+ 'snowflake' (default) to use the internal Snowflake authenticator
+ 'externalbrowser' to authenticate using your web browser and
+ Okta, ADFS or any other SAML 2.0-compliant identify provider
+ (IdP) that has been defined for your account
+ 'https://<your_okta_account_name>.okta.com' to authenticate
+ through native Okta.
+ :param warehouse: name of snowflake warehouse
+ :param database: name of snowflake database
+ :param region: name of snowflake region
+ :param role: name of snowflake role
+ :param schema: name of snowflake schema
+ :param session_parameters: You can set session-level parameters at
+ the time you connect to Snowflake
+ :param token_life_time: lifetime of the JWT Token in timedelta
+ :param token_renewal_delta: Renewal time of the JWT Token in timedelta
+ """
+
+ LIFETIME = timedelta(minutes=59) # The tokens will have a 59 minute lifetime
+ RENEWAL_DELTA = timedelta(minutes=54) # Tokens will be renewed after 54 minutes
+
+ def __init__(
+ self,
+ snowflake_conn_id: str,
+ token_life_time: timedelta = LIFETIME,
+ token_renewal_delta: timedelta = RENEWAL_DELTA,
+ *args: Any,
+ **kwargs: Any,
+ ):
+ self.snowflake_conn_id = snowflake_conn_id
+ self.token_life_time = token_life_time
+ self.token_renewal_delta = token_renewal_delta
+ super().__init__(snowflake_conn_id, *args, **kwargs)
+ self.private_key: Any = None
+
+ def get_private_key(self) -> None:
+ """Gets the private key from snowflake connection"""
+ conn = self.get_connection(self.snowflake_conn_id)
+
+ # If private_key_file is specified in the extra json, load the contents of the file as a private key.
+ # If private_key_content is specified in the extra json, use it as a private key.
+ # As a next step, specify this private key in the connection configuration.
+ # The connection password then becomes the passphrase for the private key.
+ # If your private key is not encrypted (not recommended), then leave the password empty.
+
+ private_key_file = conn.extra_dejson.get(
+ "extra__snowflake__private_key_file"
+ ) or conn.extra_dejson.get("private_key_file")
+ private_key_content = conn.extra_dejson.get(
+ "extra__snowflake__private_key_content"
+ ) or conn.extra_dejson.get("private_key_content")
+
+ private_key_pem = None
+ if private_key_content and private_key_file:
+ raise AirflowException(
+ "The private_key_file and private_key_content extra fields are mutually exclusive. "
+ "Please remove one."
+ )
+ elif private_key_file:
+ private_key_pem = Path(private_key_file).read_bytes()
+ elif private_key_content:
+ private_key_pem = private_key_content.encode()
+
+ if private_key_pem:
+ passphrase = None
+ if conn.password:
+ passphrase = conn.password.strip().encode()
+
+ self.private_key = serialization.load_pem_private_key(
+ private_key_pem, password=passphrase, backend=default_backend()
+ )
+
+ def execute_query(
+ self, sql: str, statement_count: int, query_tag: str = "", bindings: dict[str, Any] | None = None
+ ) -> list[str]:
+ """
+ Using SnowflakeSQL API, run the query in snowflake by making API request
+
+ :param sql: the sql string to be executed with possibly multiple statements
+ :param statement_count: set the MULTI_STATEMENT_COUNT field to the number of SQL statements
+ in the request
+ :param query_tag: (Optional) Query tag that you want to associate with the SQL statement.
+ For details, see https://docs.snowflake.com/en/sql-reference/parameters.html#label-query-tag
+ parameter.
+ :param bindings: (Optional) Values of bind variables in the SQL statement.
+ When executing the statement, Snowflake replaces placeholders (? and :name) in
+ the statement with these specified values.
+ """
+ conn_config = self._get_conn_params()
+
+ req_id = uuid.uuid4()
+ url = f"https://{conn_config['account']}.{conn_config['region']}.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:
+ bindings = {}
+ data = {
+ "statement": sql,
+ "resultSetMetaData": {"format": "json"},
+ "database": conn_config["database"],
+ "schema": conn_config["schema"],
+ "warehouse": conn_config["warehouse"],
+ "role": conn_config["role"],
+ "bindings": bindings,
+ "parameters": {
+ "MULTI_STATEMENT_COUNT": statement_count,
+ "query_tag": query_tag,
+ },
+ }
+ response = requests.post(url, json=data, headers=headers, params=params)
+ try:
+ response.raise_for_status()
+ except requests.exceptions.HTTPError as e: # pragma: no cover
+ raise AirflowException(
+ f"Response: {e.response.content}, " f"Status Code: {e.response.status_code}"
+ ) # pragma: no cover
Review Comment:
```suggestion
)
```
--
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
[GitHub] [airflow] eladkal commented on a diff in pull request #30698: Add SnowflakeSqlApiOperator operator
Posted by "eladkal (via GitHub)" <gi...@apache.org>.
eladkal commented on code in PR #30698:
URL: https://github.com/apache/airflow/pull/30698#discussion_r1212154907
##########
airflow/providers/snowflake/operators/snowflake.py:
##########
@@ -359,3 +368,159 @@ def __init__(
self.authenticator = authenticator
self.session_parameters = session_parameters
self.query_ids: list[str] = []
+
+
+class SnowflakeSqlApiOperator(SnowflakeOperator):
Review Comment:
This is not good.
`SnowflakeOperator` is deprecated. We should not extend it's functionality.
This class needs to be refactored around `SQLExecuteQueryOperator, `BaseOperator` or possibly have a `SnowflakeBaseOperator` class ?
--
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
[GitHub] [airflow] potiuk commented on pull request #30698: Add SnowflakeSqlApiOperator operator
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #30698:
URL: https://github.com/apache/airflow/pull/30698#issuecomment-1549445780
> @potiuk please review
I thin it's better to avoid calling specific people. I might be on holidays or focusing on something else. By calling me personally here (when I was not involved before especially) you are limiting the chances for things to be reviewed (because others might think I will do it). It's better to just ask for review `in general` - or ping people if they have been reviewing things before than asking people directly to do review..
Note to self - we need to upgrade our CODEOWNERS rules actually as they might be misleding and they are wrong in multiple cases.
--
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
[GitHub] [airflow] eladkal commented on a diff in pull request #30698: Add SnowflakeSqlApiOperator operator
Posted by "eladkal (via GitHub)" <gi...@apache.org>.
eladkal commented on code in PR #30698:
URL: https://github.com/apache/airflow/pull/30698#discussion_r1225237602
##########
airflow/providers/snowflake/operators/snowflake.py:
##########
@@ -359,3 +368,159 @@ def __init__(
self.authenticator = authenticator
self.session_parameters = session_parameters
self.query_ids: list[str] = []
+
+
+class SnowflakeSqlApiOperator(SnowflakeOperator):
Review Comment:
For reference... followup PR https://github.com/apache/airflow/pull/31751
--
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
[GitHub] [airflow] utkarsharma2 closed pull request #30698: Add SnowflakeSqlApiOperator operator
Posted by "utkarsharma2 (via GitHub)" <gi...@apache.org>.
utkarsharma2 closed pull request #30698: Add SnowflakeSqlApiOperator operator
URL: https://github.com/apache/airflow/pull/30698
--
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
[GitHub] [airflow] potiuk commented on pull request #30698: Add SnowflakeSqlApiOperator operator
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #30698:
URL: https://github.com/apache/airflow/pull/30698#issuecomment-1549477164
If we don't hear from Snowflake team, we are going t merge it soon anyway @eladkal - might be worth merging before the next wave ? But maybe we can drag attention of the Snowflake team to do more review here (there was somoene at Snowflake who looked at the changes before I recall @mik-laj ?
--
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
[GitHub] [airflow] utkarsharma2 commented on pull request #30698: Add SnowflakeSqlApiOperator operator
Posted by "utkarsharma2 (via GitHub)" <gi...@apache.org>.
utkarsharma2 commented on PR #30698:
URL: https://github.com/apache/airflow/pull/30698#issuecomment-1549024126
@phanikumv Added documents and example
--
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
[GitHub] [airflow] utkarsharma2 commented on pull request #30698: Add SnowflakeSqlApiOperator operator
Posted by "utkarsharma2 (via GitHub)" <gi...@apache.org>.
utkarsharma2 commented on PR #30698:
URL: https://github.com/apache/airflow/pull/30698#issuecomment-1549024731
@potiuk, please review
--
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
[GitHub] [airflow] eladkal commented on a diff in pull request #30698: Add SnowflakeSqlApiOperator operator
Posted by "eladkal (via GitHub)" <gi...@apache.org>.
eladkal commented on code in PR #30698:
URL: https://github.com/apache/airflow/pull/30698#discussion_r1204074285
##########
airflow/providers/snowflake/hooks/sql_api_generate_jwt.py:
##########
@@ -0,0 +1,174 @@
+# Licensed to the Apache Software Foundation (ASF) under one
Review Comment:
This is not a hook.
it looks more like util and should be in util path
##########
airflow/providers/snowflake/provider.yaml:
##########
@@ -74,6 +74,8 @@ hooks:
- integration-name: Snowflake
python-modules:
- airflow.providers.snowflake.hooks.snowflake
+ - airflow.providers.snowflake.hooks.snowflake_sql_api
+ - airflow.providers.snowflake.hooks.sql_api_generate_jwt
Review Comment:
This is not right. it's not a hook.
--
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
[GitHub] [airflow] utkarsharma2 commented on a diff in pull request #30698: Add SnowflakeSqlApiOperator operator
Posted by "utkarsharma2 (via GitHub)" <gi...@apache.org>.
utkarsharma2 commented on code in PR #30698:
URL: https://github.com/apache/airflow/pull/30698#discussion_r1205012254
##########
airflow/providers/snowflake/hooks/sql_api_generate_jwt.py:
##########
@@ -0,0 +1,174 @@
+# Licensed to the Apache Software Foundation (ASF) under one
Review Comment:
Yes, that makes sense. I'll raise a separate PR to move it. thanks for pointing it out :)
--
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
[GitHub] [airflow] utkarsharma2 commented on a diff in pull request #30698: Add SnowflakeSqlApiOperator operator
Posted by "utkarsharma2 (via GitHub)" <gi...@apache.org>.
utkarsharma2 commented on code in PR #30698:
URL: https://github.com/apache/airflow/pull/30698#discussion_r1205012254
##########
airflow/providers/snowflake/hooks/sql_api_generate_jwt.py:
##########
@@ -0,0 +1,174 @@
+# Licensed to the Apache Software Foundation (ASF) under one
Review Comment:
Yes, that makes sense. I'll move it. thanks for pointing it out :)
--
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
[GitHub] [airflow] utkarsharma2 commented on pull request #30698: Add SnowflakeSqlApiOperator operator
Posted by "utkarsharma2 (via GitHub)" <gi...@apache.org>.
utkarsharma2 commented on PR #30698:
URL: https://github.com/apache/airflow/pull/30698#issuecomment-1549465684
@potiuk Sure, I'll keep that in mind going forward. Thanks for reviewing :)
--
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
[GitHub] [airflow] eladkal commented on a diff in pull request #30698: Add SnowflakeSqlApiOperator operator
Posted by "eladkal (via GitHub)" <gi...@apache.org>.
eladkal commented on code in PR #30698:
URL: https://github.com/apache/airflow/pull/30698#discussion_r1212154907
##########
airflow/providers/snowflake/operators/snowflake.py:
##########
@@ -359,3 +368,159 @@ def __init__(
self.authenticator = authenticator
self.session_parameters = session_parameters
self.query_ids: list[str] = []
+
+
+class SnowflakeSqlApiOperator(SnowflakeOperator):
Review Comment:
This is not good.
`SnowflakeOperator` is deprecated. We should not extend it's functionality.
This class needs to be refactored around `SQLExecuteQueryOperator` or `BaseOperator` / possibly have a `SnowflakeBaseOperator` class ?
--
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
[GitHub] [airflow] eladkal commented on a diff in pull request #30698: Add SnowflakeSqlApiOperator operator
Posted by "eladkal (via GitHub)" <gi...@apache.org>.
eladkal commented on code in PR #30698:
URL: https://github.com/apache/airflow/pull/30698#discussion_r1212154907
##########
airflow/providers/snowflake/operators/snowflake.py:
##########
@@ -359,3 +368,159 @@ def __init__(
self.authenticator = authenticator
self.session_parameters = session_parameters
self.query_ids: list[str] = []
+
+
+class SnowflakeSqlApiOperator(SnowflakeOperator):
Review Comment:
This is not good.
`SnowflakeOperator` is deprecated. We should not extend it's functionality.
This class needs to be refactored around `SQLExecuteQueryOperator`, `BaseOperator` or possibly have a `SnowflakeBaseOperator` class ?
--
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