You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by GitBox <gi...@apache.org> on 2021/08/30 10:47:15 UTC

[GitHub] [airflow] JavierLopezT commented on a change in pull request #17741: Add Snowflake DQ Operators

JavierLopezT commented on a change in pull request #17741:
URL: https://github.com/apache/airflow/pull/17741#discussion_r698382156



##########
File path: airflow/providers/snowflake/operators/snowflake.py
##########
@@ -125,3 +126,285 @@ def execute(self, context: Any) -> None:
 
         if self.do_xcom_push:
             return execution_info
+
+
+class _SnowflakeDbHookMixin:
+    def get_db_hook(self) -> SnowflakeHook:
+        """
+        Create and return SnowflakeHook.
+
+        :return: a SnowflakeHook instance.
+        :rtype: SnowflakeHook
+        """
+        return SnowflakeHook(
+            snowflake_conn_id=self.snowflake_conn_id,
+            warehouse=self.warehouse,
+            database=self.database,
+            role=self.role,
+            schema=self.schema,
+            authenticator=self.authenticator,
+            session_parameters=self.session_parameters,
+        )
+
+
+class SnowflakeCheckOperator(_SnowflakeDbHookMixin, SQLCheckOperator):
+    """
+    Performs a check against Snowflake. The ``SnowflakeCheckOperator`` expects
+    a sql query that will return a single row. Each value on that
+    first row is evaluated using python ``bool`` casting. If any of the
+    values return ``False`` the check is failed and errors out.
+
+    Note that Python bool casting evals the following as ``False``:
+
+    * ``False``
+    * ``0``
+    * Empty string (``""``)
+    * Empty list (``[]``)
+    * Empty dictionary or set (``{}``)
+
+    Given a query like ``SELECT COUNT(*) FROM foo``, it will fail only if
+    the count ``== 0``. You can craft much more complex query that could,
+    for instance, check that the table has the same number of rows as
+    the source table upstream, or that the count of today's partition is
+    greater than yesterday's partition, or that a set of metrics are less
+    than 3 standard deviation for the 7 day average.
+
+    This operator can be used as a data quality check in your pipeline, and
+    depending on where you put it in your DAG, you have the choice to
+    stop the critical path, preventing from
+    publishing dubious data, or on the side and receive email alerts
+    without stopping the progress of the DAG.
+
+    :param sql: the sql code to be executed. (templated)
+    :type sql: Can receive a str representing a sql statement,
+        a list of str (sql statements), or reference to a template file.
+        Template reference are recognized by str ending in '.sql'
+    :param snowflake_conn_id: Reference to
+        :ref:`Snowflake connection id<howto/connection:snowflake>`
+    :type snowflake_conn_id: str
+    :param autocommit: if True, each command is automatically committed.
+        (default value: True)
+    :type autocommit: bool
+    :param parameters: (optional) the parameters to render the SQL query with.
+    :type parameters: dict or iterable
+    :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
+    """
+
+    template_fields = ('sql',)
+    template_ext = ('.sql',)
+    ui_color = '#ededed'
+
+    def __init__(
+        self,
+        *,
+        sql: 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, **kwargs)
+        self.snowflake_conn_id = snowflake_conn_id
+        self.sql = sql
+        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 = []
+
+
+class SnowflakeValueCheckOperator(_SnowflakeDbHookMixin, SQLValueCheckOperator):
+    """
+    Performs a simple value check using sql code.

Review comment:
       I know that this is copied from https://github.com/apache/airflow/blob/main/airflow/operators/sql.py#L149, but I think it's unclear what the operator really is. I think adding something like 'it checks against a value that you define' would be really useful

##########
File path: airflow/providers/snowflake/operators/snowflake.py
##########
@@ -125,3 +126,285 @@ def execute(self, context: Any) -> None:
 
         if self.do_xcom_push:
             return execution_info
+
+
+class _SnowflakeDbHookMixin:
+    def get_db_hook(self) -> SnowflakeHook:
+        """
+        Create and return SnowflakeHook.
+
+        :return: a SnowflakeHook instance.
+        :rtype: SnowflakeHook
+        """
+        return SnowflakeHook(
+            snowflake_conn_id=self.snowflake_conn_id,
+            warehouse=self.warehouse,
+            database=self.database,
+            role=self.role,
+            schema=self.schema,
+            authenticator=self.authenticator,
+            session_parameters=self.session_parameters,
+        )
+
+
+class SnowflakeCheckOperator(_SnowflakeDbHookMixin, SQLCheckOperator):
+    """
+    Performs a check against Snowflake. The ``SnowflakeCheckOperator`` expects
+    a sql query that will return a single row. Each value on that
+    first row is evaluated using python ``bool`` casting. If any of the
+    values return ``False`` the check is failed and errors out.
+
+    Note that Python bool casting evals the following as ``False``:
+
+    * ``False``
+    * ``0``
+    * Empty string (``""``)
+    * Empty list (``[]``)
+    * Empty dictionary or set (``{}``)
+
+    Given a query like ``SELECT COUNT(*) FROM foo``, it will fail only if
+    the count ``== 0``. You can craft much more complex query that could,
+    for instance, check that the table has the same number of rows as
+    the source table upstream, or that the count of today's partition is
+    greater than yesterday's partition, or that a set of metrics are less
+    than 3 standard deviation for the 7 day average.
+
+    This operator can be used as a data quality check in your pipeline, and
+    depending on where you put it in your DAG, you have the choice to
+    stop the critical path, preventing from
+    publishing dubious data, or on the side and receive email alerts
+    without stopping the progress of the DAG.
+
+    :param sql: the sql code to be executed. (templated)
+    :type sql: Can receive a str representing a sql statement,
+        a list of str (sql statements), or reference to a template file.
+        Template reference are recognized by str ending in '.sql'
+    :param snowflake_conn_id: Reference to
+        :ref:`Snowflake connection id<howto/connection:snowflake>`
+    :type snowflake_conn_id: str
+    :param autocommit: if True, each command is automatically committed.

Review comment:
       Given that the the sqls are meant to be SELECT statements, does it make sense to have this param?

##########
File path: airflow/providers/snowflake/operators/snowflake.py
##########
@@ -125,3 +126,285 @@ def execute(self, context: Any) -> None:
 
         if self.do_xcom_push:
             return execution_info
+
+
+class _SnowflakeDbHookMixin:

Review comment:
       Why a class and not just a method inside the SnowflakeCheckOperator?

##########
File path: airflow/providers/snowflake/operators/snowflake.py
##########
@@ -125,3 +126,285 @@ def execute(self, context: Any) -> None:
 
         if self.do_xcom_push:
             return execution_info
+
+
+class _SnowflakeDbHookMixin:
+    def get_db_hook(self) -> SnowflakeHook:
+        """
+        Create and return SnowflakeHook.
+
+        :return: a SnowflakeHook instance.
+        :rtype: SnowflakeHook
+        """
+        return SnowflakeHook(
+            snowflake_conn_id=self.snowflake_conn_id,
+            warehouse=self.warehouse,
+            database=self.database,
+            role=self.role,
+            schema=self.schema,
+            authenticator=self.authenticator,
+            session_parameters=self.session_parameters,
+        )
+
+
+class SnowflakeCheckOperator(_SnowflakeDbHookMixin, SQLCheckOperator):
+    """
+    Performs a check against Snowflake. The ``SnowflakeCheckOperator`` expects
+    a sql query that will return a single row. Each value on that
+    first row is evaluated using python ``bool`` casting. If any of the
+    values return ``False`` the check is failed and errors out.
+
+    Note that Python bool casting evals the following as ``False``:
+
+    * ``False``
+    * ``0``
+    * Empty string (``""``)
+    * Empty list (``[]``)
+    * Empty dictionary or set (``{}``)
+
+    Given a query like ``SELECT COUNT(*) FROM foo``, it will fail only if
+    the count ``== 0``. You can craft much more complex query that could,
+    for instance, check that the table has the same number of rows as
+    the source table upstream, or that the count of today's partition is
+    greater than yesterday's partition, or that a set of metrics are less
+    than 3 standard deviation for the 7 day average.
+
+    This operator can be used as a data quality check in your pipeline, and
+    depending on where you put it in your DAG, you have the choice to
+    stop the critical path, preventing from
+    publishing dubious data, or on the side and receive email alerts
+    without stopping the progress of the DAG.
+
+    :param sql: the sql code to be executed. (templated)
+    :type sql: Can receive a str representing a sql statement,
+        a list of str (sql statements), or reference to a template file.
+        Template reference are recognized by str ending in '.sql'
+    :param snowflake_conn_id: Reference to
+        :ref:`Snowflake connection id<howto/connection:snowflake>`
+    :type snowflake_conn_id: str
+    :param autocommit: if True, each command is automatically committed.
+        (default value: True)
+    :type autocommit: bool
+    :param parameters: (optional) the parameters to render the SQL query with.
+    :type parameters: dict or iterable
+    :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
+    """
+
+    template_fields = ('sql',)
+    template_ext = ('.sql',)
+    ui_color = '#ededed'
+
+    def __init__(
+        self,
+        *,
+        sql: 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, **kwargs)
+        self.snowflake_conn_id = snowflake_conn_id
+        self.sql = sql
+        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 = []
+
+
+class SnowflakeValueCheckOperator(_SnowflakeDbHookMixin, SQLValueCheckOperator):
+    """
+    Performs a simple value check using sql code.
+
+    :param sql: the sql to be executed
+    :type sql: str
+    :param snowflake_conn_id: Reference to
+        :ref:`Snowflake connection id<howto/connection:snowflake>`
+    :type snowflake_conn_id: str
+    :param autocommit: if True, each command is automatically committed.
+        (default value: True)
+    :type autocommit: bool
+    :param parameters: (optional) the parameters to render the SQL query with.
+    :type parameters: dict or iterable
+    :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: Any,

Review comment:
       Is it really Any? Or union str and list?

##########
File path: airflow/providers/snowflake/operators/snowflake.py
##########
@@ -125,3 +126,285 @@ def execute(self, context: Any) -> None:
 
         if self.do_xcom_push:
             return execution_info
+
+
+class _SnowflakeDbHookMixin:
+    def get_db_hook(self) -> SnowflakeHook:
+        """
+        Create and return SnowflakeHook.
+
+        :return: a SnowflakeHook instance.
+        :rtype: SnowflakeHook
+        """
+        return SnowflakeHook(
+            snowflake_conn_id=self.snowflake_conn_id,
+            warehouse=self.warehouse,
+            database=self.database,
+            role=self.role,
+            schema=self.schema,
+            authenticator=self.authenticator,
+            session_parameters=self.session_parameters,
+        )
+
+
+class SnowflakeCheckOperator(_SnowflakeDbHookMixin, SQLCheckOperator):
+    """
+    Performs a check against Snowflake. The ``SnowflakeCheckOperator`` expects
+    a sql query that will return a single row. Each value on that
+    first row is evaluated using python ``bool`` casting. If any of the
+    values return ``False`` the check is failed and errors out.
+
+    Note that Python bool casting evals the following as ``False``:
+
+    * ``False``
+    * ``0``
+    * Empty string (``""``)
+    * Empty list (``[]``)
+    * Empty dictionary or set (``{}``)
+
+    Given a query like ``SELECT COUNT(*) FROM foo``, it will fail only if
+    the count ``== 0``. You can craft much more complex query that could,
+    for instance, check that the table has the same number of rows as
+    the source table upstream, or that the count of today's partition is
+    greater than yesterday's partition, or that a set of metrics are less
+    than 3 standard deviation for the 7 day average.
+
+    This operator can be used as a data quality check in your pipeline, and
+    depending on where you put it in your DAG, you have the choice to
+    stop the critical path, preventing from
+    publishing dubious data, or on the side and receive email alerts
+    without stopping the progress of the DAG.
+
+    :param sql: the sql code to be executed. (templated)
+    :type sql: Can receive a str representing a sql statement,
+        a list of str (sql statements), or reference to a template file.
+        Template reference are recognized by str ending in '.sql'
+    :param snowflake_conn_id: Reference to
+        :ref:`Snowflake connection id<howto/connection:snowflake>`
+    :type snowflake_conn_id: str
+    :param autocommit: if True, each command is automatically committed.
+        (default value: True)
+    :type autocommit: bool
+    :param parameters: (optional) the parameters to render the SQL query with.
+    :type parameters: dict or iterable
+    :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
+    """
+
+    template_fields = ('sql',)
+    template_ext = ('.sql',)
+    ui_color = '#ededed'
+
+    def __init__(
+        self,
+        *,
+        sql: 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, **kwargs)
+        self.snowflake_conn_id = snowflake_conn_id
+        self.sql = sql
+        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 = []
+
+
+class SnowflakeValueCheckOperator(_SnowflakeDbHookMixin, SQLValueCheckOperator):
+    """
+    Performs a simple value check using sql code.
+
+    :param sql: the sql to be executed
+    :type sql: str
+    :param snowflake_conn_id: Reference to
+        :ref:`Snowflake connection id<howto/connection:snowflake>`
+    :type snowflake_conn_id: str
+    :param autocommit: if True, each command is automatically committed.
+        (default value: True)
+    :type autocommit: bool
+    :param parameters: (optional) the parameters to render the SQL query with.
+    :type parameters: dict or iterable
+    :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: Any,
+        pass_value: Any,
+        tolerance: Any = None,
+        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, pass_value=pass_value, tolerance=tolerance, **kwargs)
+        self.snowflake_conn_id = snowflake_conn_id
+        self.sql = sql
+        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 = []
+
+
+class SnowflakeIntervalCheckOperator(_SnowflakeDbHookMixin, SQLIntervalCheckOperator):
+    """
+    Checks that the values of metrics given as SQL expressions are within
+    a certain tolerance of the ones from days_back before.
+
+    This method constructs a query like so ::
+
+        SELECT {metrics_threshold_dict_key} FROM {table}
+        WHERE {date_filter_column}=<date>
+
+    :param table: the table name
+    :type table: str
+    :param days_back: number of days between ds and the ds we want to check

Review comment:
       If this is here, I guess `tolerance` and `pass_value` should be in the docstring of `SnowflakeValueCheckOperator`. It seems that for Sphinx is not necessary https://stackoverflow.com/a/2025599/9621172 but certainly it really helps to have them here if you are just reading the code




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