How to use flowetl - 10 common examples

# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.

from typing import Dict, Optional

from airflow.operators.postgres_operator import PostgresOperator
from flowetl.mixins.table_name_macros_mixin import TableNameMacrosMixin


class CreateForeignStagingTableOperator(TableNameMacrosMixin, PostgresOperator):
    def __init__(
        self,
        *,
        filename: str,
        fields: Dict[str, str],
        program: Optional[str] = None,
        header: bool = True,
        delimiter: str = ",",
        quote: str = '"',
        escape: str = '"',
        null: str = "",
        encoding: Optional[str] = None,
        **kwargs,
    ) -> None:
        """
        Operator which uses file_fdw to create a table which can be used to read a flat file.

class_name : str
        Class name for the operator
    sql : str
        SQL string to use as wrapper

    Returns
    -------
    Type

    """
    from flowetl.mixins.table_name_macros_mixin import TableNameMacrosMixin
    from airflow.operators.postgres_operator import PostgresOperator

    return type(
        class_name,
        (TableNameMacrosMixin, WrappingSQLMixin, PostgresOperator),
        dict(wrapper_sql=sql),
    )

# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
from airflow.operators.sensors import SqlSensor
from airflow.utils.decorators import apply_defaults
from flowetl.mixins.table_name_macros_mixin import TableNameMacrosMixin


class FileFluxSensor(TableNameMacrosMixin, SqlSensor):
    """
    The file flux sensor monitors a file for a short time to check if it is still
    being modified.

    Parameters
    ----------
    conn_id : str
        Connection to use
    flux_check_interval : int
        Number of seconds to wait between checks that a file is stable
    filename : str
        jinja templated string providing the path to the file to check
    kwargs : dict
        Passed to airflow.operators.sensors.SqlSensor
    See Also
    --------

Returns
    -------
    Type
        New operator class

    """
    from flowetl.mixins.table_name_macros_mixin import TableNameMacrosMixin

    if is_sensor:
        from airflow.sensors.sql_sensor import SqlSensor as op_base
    else:
        from airflow.operators.postgres_operator import PostgresOperator as op_base

    return type(
        class_name, (TableNameMacrosMixin, FixedSQLMixin, op_base), dict(fixed_sql=sql),
    )

-------
        Type
            New operator class

        """
    from flowetl.mixins.fixed_sql_mixin import FixedSQLMixin
    from flowetl.mixins.table_name_macros_mixin import TableNameMacrosMixin

    if is_sensor:
        from airflow.sensors.sql_sensor import SqlSensor as op_base
    else:
        from airflow.operators.postgres_operator import PostgresOperator as op_base

    return type(
        class_name,
        (TableNameMacrosMixin, ParamsMixin, FixedSQLMixin, op_base),
        dict(fixed_sql=sql, named_params=params),
    )

# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
from airflow.operators.postgres_operator import PostgresOperator
from flowetl.mixins.table_name_macros_mixin import TableNameMacrosMixin


class AnalyzeOperator(TableNameMacrosMixin, PostgresOperator):
    """
    The analyze operator triggers the postgres analyze command on a table.

    Parameters
    ----------
    target : str
        jinja templated schema qualified table name.
    kwargs : dict
        Passed to airflow.operators.postgres_operator.PostgresOperator
    """

    def __init__(self, *, target: str, **kwargs) -> None:
        super().__init__(
            sql=f"ANALYZE {target};", **kwargs
        )  # Need an f-string to let us use templating with the target

# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.

from flowetl.mixins.fixed_sql_mixin import fixed_sql_operator

UpdateETLTableOperator = fixed_sql_operator(
    class_name="UpdateETLTableOperator",
    sql="""
        INSERT INTO etl.etl_records (cdr_type, cdr_date, state, timestamp) VALUES ('{{ params.cdr_type }}', '{{ ds }}'::DATE, 'ingested', NOW());

# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.

from flowetl.mixins.fixed_sql_mixin import fixed_sql_operator

DataPresentSensor = fixed_sql_operator(
    class_name="DataPresentSensor",
    sql="SELECT * FROM {{ staging_table }} LIMIT 1;",
    is_sensor=True,
)

# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.

from flowetl.mixins.fixed_sql_mixin import fixed_sql_operator

AddConstraintsOperator = fixed_sql_operator(
    class_name="AddConstraintsOperator",
    sql="""
        ALTER TABLE {{ extract_table }}

# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.

from flowetl.mixins.fixed_sql_mixin import fixed_sql_operator

AttachOperator = fixed_sql_operator(
    class_name="AttachOperator",
    sql="""
        DROP TABLE IF EXISTS {{ final_table }};