[GitHub] [airflow] denimalpaca commented on a change in pull request #18413: Add/Update Snowflake SQL Operators

[GitBox <gi...@apache.org>]

denimalpaca commented on a change in pull request #18413:
URL: https://github.com/apache/airflow/pull/18413#discussion_r715708316



##########
File path: airflow/providers/snowflake/operators/snowflake.py
##########
@@ -409,3 +430,175 @@ def __init__(
 
     def get_db_hook(self) -> SnowflakeHook:
         return get_db_hook(self)
+
+
+class SnowflakeThresholdCheckOperator(SQLThresholdCheckOperator):
+    """
+    Performs a value check using sql code against a minimum threshold
+    and a maximum threshold. Thresholds can be in the form of a numeric
+    value OR a sql statement that results a numeric.
+
+    :param sql: the sql to be executed. (templated)
+    :type sql: str
+    :param snowflake_conn_id: Reference to
+        :ref:`Snowflake connection id<howto/connection:snowflake>`
+    :type snowflake_conn_id: str
+    :param min_threshold: numerical value or min threshold sql to be executed (templated)
+    :type min_threshold: numeric or str
+    :param max_threshold: numerical value or max threshold sql to be executed (templated)
+    :type max_threshold: numeric or str
+    :param autocommit: if True, each command is automatically committed.
+        (default value: True)
+    :type autocommit: bool
+    :param warehouse: name of warehouse (will overwrite any warehouse
+        defined in the connection's extra JSON)
+    :type warehouse: str
+    :param database: name of database (will overwrite database defined
+        in connection)
+    :type database: str
+    :param schema: name of schema (will overwrite schema defined in
+        connection)
+    :type schema: str
+    :param role: name of role (will overwrite any role defined in
+        connection's extra JSON)
+    :type role: str
+    :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.
+    :type authenticator: str
+    :param session_parameters: You can set session-level parameters at
+        the time you connect to Snowflake
+    :type session_parameters: dict
+    """
+
+    def __init__(
+        self,
+        *,
+        sql: str,
+        min_threshold: Any,
+        max_threshold: Any,
+        snowflake_conn_id: str = 'snowflake_default',
+        parameters: Optional[dict] = None,
+        autocommit: bool = True,
+        do_xcom_push: bool = True,
+        warehouse: Optional[str] = None,
+        database: Optional[str] = None,
+        role: Optional[str] = None,
+        schema: Optional[str] = None,
+        authenticator: Optional[str] = None,
+        session_parameters: Optional[dict] = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(
+            sql=sql,
+            min_threshold=min_threshold,
+            max_threshold=max_threshold,
+            **kwargs,
+        )
+
+        self.snowflake_conn_id = snowflake_conn_id
+        self.autocommit = autocommit
+        self.do_xcom_push = do_xcom_push
+        self.parameters = parameters
+        self.warehouse = warehouse
+        self.database = database
+        self.role = role
+        self.schema = schema
+        self.authenticator = authenticator
+        self.session_parameters = session_parameters
+        self.query_ids = []
+
+    def get_db_hook(self) -> SnowflakeHook:
+        return get_db_hook(self)
+
+
+class BranchSnowflakeOperator(BranchSQLOperator):

Review comment:
       I like this solution a lot, would also handle all the other cases (BigQuery specifically, which this and my last PR was based on). Seems like it would make more sense for me to open a separate PR for that solution, though.




-- 
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