flowtask 5.8.4__cp39-cp39-manylinux2014_x86_64.manylinux_2_17_x86_64.whl

flowtask/components/TableDelete.py

@@ -0,0 +1,312 @@
+import asyncio
+
+# logging system
+import logging
+import datetime
+from decimal import Decimal
+from typing import Any
+from collections.abc import Callable
+import pandas as pd
+import numpy as np
+from asyncdb.exceptions import NoDataFound
+from asyncdb.models import Model
+from asyncdb.drivers.pg import pg
+
+# for database
+from querysource.conf import (
+    sqlalchemy_url,
+    default_dsn,
+    DB_STATEMENT_TIMEOUT,
+    DB_SESSION_TIMEOUT,
+    DB_KEEPALIVE_IDLE,
+)
+from sqlalchemy.orm import Session
+from sqlalchemy import create_engine, text
+from sqlalchemy.pool import NullPool
+from ..exceptions import ComponentError, DataNotFound
+from ..utils import AttrDict
+from .flow import FlowComponent
+from ..interfaces.qs import QSSupport
+
+
+dtypes = {
+    "varchar": str,
+    "string": str,
+    "object": str,
+    "int": int,
+    "int4": int,
+    "int64": int,
+    "uint64": int,
+    "Int64Dtype": np.int64,
+    "Int64": int,
+    "Int32": int,
+    "Int16": int,
+    "Int8": int,
+    "float64": Decimal,
+    "float": Decimal,
+    "bool": bool,
+    "datetime64[ns]": datetime.datetime,
+    "datetime64[ns, UTC]": datetime.datetime,
+}
+
+
+class TableDelete(QSSupport, FlowComponent):
+    """
+    TableDelete
+
+    Overview
+
+        The TableDelete class is a component for deleting rows from a SQL table based on specified primary keys. It uses
+        Pandas DataFrames to identify rows to delete and performs the deletion using SQLAlchemy and asyncpg.
+
+    .. table:: Properties
+    :widths: auto
+
+        +------------------+----------+-----------+--------------------------------------------------------------------------------------+
+        | Name             | Required | Description                                                                                      |
+        +------------------+----------+-----------+--------------------------------------------------------------------------------------+
+        | flavor           |   No     | The database flavor, defaults to "postgresql".                                                    |
+        +------------------+----------+-----------+--------------------------------------------------------------------------------------+
+        | pk               |   Yes    | A list of primary key columns to identify rows for deletion.                                      |
+        +------------------+----------+-----------+--------------------------------------------------------------------------------------+
+        | data             |   Yes    | The data containing rows to be deleted, can be a Pandas DataFrame or other supported format.      |
+        +------------------+----------+-----------+--------------------------------------------------------------------------------------+
+        | _engine          |   No     | The SQLAlchemy engine for database connections.                                                   |
+        +------------------+----------+-----------+--------------------------------------------------------------------------------------+
+        | tablename        |   Yes    | The name of the table from which rows will be deleted.                                            |
+        +------------------+----------+-----------+--------------------------------------------------------------------------------------+
+        | schema           |   No     | The schema name, defaults to an empty string.                                                     |
+        +------------------+----------+-----------+--------------------------------------------------------------------------------------+
+        | _dsn             |   No     | The Data Source Name for the database connection.                                                 |
+        +------------------+----------+-----------+--------------------------------------------------------------------------------------+
+        | multi            |   No     | A flag indicating if multiple DataFrame deletions are supported, defaults to False.               |
+        +------------------+----------+-----------+--------------------------------------------------------------------------------------+
+
+    Return
+
+        The methods in this class manage the deletion of rows from a SQL table based on primary keys, including initialization,
+        execution, and result handling. It ensures proper handling of database connections and provides metrics on the deleted rows.
+
+    
+
+        Example:
+
+        ```yaml
+        TableDelete:
+          tablename: totem_heartbeats
+          flavor: postgres
+          schema: banco_chile
+          pk:
+          - hearbeat_id
+        ```
+
+    """  # noqa
+
+    flavor: str = "postgresql"
+
+    def __init__(
+        self,
+        loop: asyncio.AbstractEventLoop = None,
+        job: Callable = None,
+        stat: Callable = None,
+        **kwargs,
+    ):
+        self.pk = []
+        self.data: Any = None
+        self._engine = None
+        self.tablename = ""
+        self.schema = ""
+        dsn = kwargs.get('dsn', '')
+        if dsn and dsn.startswith('postgres:'):
+            dsn = dsn.replace('postgres:', 'postgresql:')
+        self._dsn = self.get_config_value(dsn, dsn)
+        self.multi = bool(kwargs.pop('multi', False))
+        self._driver: str = kwargs.pop('driver', 'pg')
+        super(TableDelete, self).__init__(
+            loop=loop,
+            job=job,
+            stat=stat,
+            **kwargs
+        )
+
+    def get_engine(self, dsn: str = None):
+        if not dsn:
+            dsn = sqlalchemy_url
+        try:
+            self._engine = create_engine(dsn, echo=False, poolclass=NullPool)
+            self._session = Session(bind=self._engine)
+        except Exception as err:
+            logging.exception(err)
+
+    async def start(self, **kwargs):
+        """Get Pandas Dataframe."""
+        self.data = None
+        if self.previous:
+            self.data = self.input
+        else:
+            raise NoDataFound("TableDelete: Data missing")
+        await super().start(**kwargs)
+        try:
+            # getting sqlAlchemy engine
+            self.get_engine(dsn=self._dsn)
+        except Exception as err:
+            raise ComponentError(f"Connection Error: {err}") from err
+        if self.data is None:
+            raise NoDataFound("TableDelete: Data missing")
+        elif isinstance(self.data, pd.DataFrame):
+            # detect data type for colums
+            columns = list(self.data.columns)
+            for column in columns:
+                t = self.data[column].dtype
+                if isinstance(t, pd.core.dtypes.dtypes.DatetimeTZDtype):
+                    self.data[column] = pd.to_datetime(
+                        self.data[column],
+                        format="%Y-%m-%dT%H:%M:%S.%f%z",
+                        cache=True,
+                        utc=True,
+                    )
+                    self.data[column].dt.tz_convert("UTC")
+                elif str(t) == "datetime64[ns]":
+                    tmp_data = self.data.copy()
+                    tmp_data[column] = pd.to_datetime(
+                        self.data[column],
+                        format="%Y-%m-%dT%H:%M:%S.%f%z",
+                        cache=True,
+                        utc=True,
+                    )
+                    self.data = tmp_data.copy()
+        elif self.multi is True:
+            # iteration over every Pandas DT:
+            try:
+                result = self.data.items()
+            except Exception as err:
+                raise ComponentError(
+                    "Invalid Result type for Attribute Multiple: {err}"
+                ) from err
+            for name, rs in result:
+                el = getattr(self, name)
+                if not isinstance(rs, pd.DataFrame):
+                    raise ComponentError("Invalid Resultset: not a Dataframe")
+        else:
+            # incompatible datasource:
+            raise DataNotFound("TableDelete: Incompatible Data Source")
+
+    async def close(self):
+        """Closing Operations."""
+        if self._engine:
+            try:
+                self._engine.dispose()
+            except Exception:
+                pass
+
+    async def table_delete(self, elem, df):
+        """Running the process of Upsert-delete."""
+        pk = elem.pk
+        options = {
+            "if_exists": "append",
+            "index_label": pk,
+            "index": False,
+            "method": "multi",
+        }
+        if hasattr(elem, "sql_options"):
+            options = {**options, **elem.sql_options}
+        tablename = elem.tablename
+        try:
+            schema_name = elem.schema
+        except AttributeError:
+            schema_name = "public"
+        options["schema"] = schema_name
+        cols = []
+        for field in pk:
+            datatype = df.dtypes[field]
+
+            t = dtypes[f"{datatype!s}"]
+            f = (field, t)
+            cols.append(f)
+        # make a intermediate model:
+        try:
+            cls = Model.make_model(
+                name=f"{tablename!s}_deleted", schema=schema_name, fields=cols
+            )
+        except Exception as err:
+            logging.exception(
+                f"TableDelete: Error creating DB Model for {tablename}: {err}"
+            )
+            return False
+        result = None
+        if sql := cls.model(dialect="sql"):
+            print("SQL is: ", sql)
+            deltable = text(
+                f"DROP TABLE IF EXISTS {schema_name}.{tablename!s}_deleted;"
+            )
+            qry = text(sql)
+            try:
+                with self._engine.begin() as connection:
+                    connection.execute(deltable)
+                    result = connection.execute(qry)
+                logging.debug(f"Created Temp Table: {cls!s}")
+            except Exception as error:
+                raise ComponentError(f"Error on Table creation: {error}") from error
+        else:
+            raise ComponentError("Cannot Create TEMP Table")
+
+        # table exists: copy data into table:
+        new = df[pk].copy()
+        with self._engine.begin() as connection:
+            new.to_sql(f"{tablename!s}_deleted", con=connection, **options)
+        # data is on temp Table, deleted
+        columns = ", ".join(pk)
+        # TODO: using an EXISTS instead IN to speed-up delete
+        delete = f"""DELETE FROM {schema_name!s}.{tablename!s} WHERE ({columns})
+        IN (SELECT {columns} FROM {schema_name}.{tablename!s}_deleted)"""
+        print(delete)
+        connection = await self.create_connection()
+        try:
+            async with await connection.connection() as conn:
+                try:
+                    result, error = await conn.execute(sentence=delete)
+                    print("TableDelete: ", result)
+                    self.add_metric("TABLE_DELETED", result)
+                    if error:
+                        raise ComponentError(
+                            f"Error Deleting Data: {error}"
+                        )
+                    if error:
+                        raise ComponentError(
+                            f"Error on TableDelete: {error}"
+                        )
+                except Exception as err:
+                    logging.error(err)
+                finally:
+                    # at now, delete the table:
+                    drop = f"DROP TABLE IF EXISTS {schema_name}.{tablename!s}_deleted"
+                    result, error = await conn.execute(sentence=drop)
+        finally:
+            await connection.close()
+
+    async def run(self):
+        """Run TableDelete."""
+        if self.multi is False:
+            # set the number of rows:
+            numrows = len(self.data.index)
+            self._variables[f"{self.StepName}_NUMROWS"] = numrows
+            await self.table_delete(self, self.data)
+            self._result = self.data
+            return True
+        else:
+            # running in multi Mode
+            try:
+                result = self.data.items()
+            except Exception as err:
+                raise ComponentError(
+                    "Invalid Result type for Attribute Multiple: {err}"
+                ) from err
+            for name, rs in result:
+                try:
+                    el = getattr(self, name)
+                except AttributeError:
+                    continue
+                await self.table_delete(AttrDict(el), rs)
+            # return the same dataframe
+            self._result = self.data
+            return True

flowtask/components/TableInput.py

@@ -0,0 +1,143 @@
+from decimal import Decimal
+import datetime
+import dask.dataframe as dd
+import pandas
+import numpy as np
+from ..exceptions import DataNotFound
+from .TableBase import TableBase
+
+dtypes = {
+    "varchar": str,
+    "string": str,
+    "object": str,
+    "int": int,
+    "int4": int,
+    "int64": np.int64,
+    "uint64": np.int64,
+    "Int64Dtype": np.int64,
+    "Int64": np.int64,
+    "Int8": int,
+    "float64": Decimal,
+    "float": Decimal,
+    "bool": bool,
+    "datetime64[ns]": datetime.datetime,
+    "datetime64[ns, UTC]": datetime.datetime,
+}
+
+
+class TableInput(TableBase):
+    """
+    TableInput
+
+        Overview
+
+            The TableInput class is a component for loading data from a SQL table or custom SQL query into a Pandas DataFrame.
+            It supports large tables by utilizing Dask for partitioned loading, along with options for column trimming,
+            primary key setting, and removal of empty rows or columns.
+
+        .. table:: Properties
+        :widths: auto
+
+            +----------------+----------+-----------+-------------------------------------------------------------------+
+            | Name           | Required | Summary                                                                       |
+            +----------------+----------+-----------+-------------------------------------------------------------------+
+            | tablename      |   Yes    | The name of the SQL table to load into a DataFrame.                           |
+            +----------------+----------+-----------+-------------------------------------------------------------------+
+            | schema         |   No     | Database schema where the table is located, if applicable.                    |
+            +----------------+----------+-----------+-------------------------------------------------------------------+
+            | query          |   No     | SQL query string to load data if `tablename` is not used.                     |
+            +----------------+----------+-----------+-------------------------------------------------------------------+
+            | chunksize      |   No     | Number of rows per chunk when reading data in chunks.                         |
+            +----------------+----------+-----------+-------------------------------------------------------------------+
+            | bigfile        |   No     | Boolean indicating if Dask should be used for large tables.                   |
+            +----------------+----------+-----------+-------------------------------------------------------------------+
+            | drop_empty     |   No     | Boolean to drop columns and rows with only null values.                       |
+            +----------------+----------+-----------+-------------------------------------------------------------------+
+            | trim           |   No     | Boolean to trim whitespace from string columns.                               |
+            +----------------+----------+-----------+-------------------------------------------------------------------+
+            | pk             |   No     | Dictionary specifying primary key columns and settings for DataFrame indexing.|
+            +----------------+----------+-----------+-------------------------------------------------------------------+
+
+        Returns
+
+            This component returns a DataFrame populated with data from the specified SQL table or query. If `drop_empty` is set,
+            columns and rows with only null values are removed. Metrics such as row and column counts are recorded, along with
+            information on the table or query used for data extraction. If no data is found, a `DataNotFound` exception is raised.
+    
+
+        Example:
+
+        ```yaml
+        TableInput:
+          tablename: sales_raw
+          schema: epson
+          chunksize: 1000
+          pk:
+            columns:
+            - sales_id
+        ```
+
+    """
+
+
+    async def run(self):
+        df = None
+        if hasattr(self, "tablename"):  # getting a direct Table
+            # run table to pandas
+            if hasattr(self, "bigfile"):
+                dk = dd.read_sql_table(
+                    self.tablename,
+                    self._engine,
+                    npartitions=10,
+                    head_rows=100,
+                    **self.params,
+                )
+                df = dk.compute()
+            else:
+                tp = pandas.read_sql_table(
+                    self.tablename,
+                    self._engine,
+                    schema=self.schema,
+                    chunksize=self.chunksize,
+                    **self.params,
+                )
+                df = pandas.concat(tp, ignore_index=True)
+        elif self.query:
+            # run the query to pandas
+            tp = pandas.read_sql_query(
+                self.query, self._engine, chunksize=self.chunksize, **self.params
+            )
+            df = pandas.concat(tp, ignore_index=True)
+        else:
+            return False
+        if df is None or df.empty:
+            raise DataNotFound("TableInput: Empty Dataframe")
+        # removing empty cols
+        if hasattr(self, "drop_empty"):
+            df.dropna(axis=1, how="all", inplace=True)
+            df.dropna(axis=0, how="all", inplace=True)
+        if hasattr(self, "trim"):
+            u = df.select_dtypes(include=["object", "string"])
+            df[u.columns] = df[u.columns].astype(str).str.strip()
+        # define the primary keys for DataFrame
+        if hasattr(self, "pk"):
+            try:
+                columns = self.pk["columns"]
+                del self.pk["columns"]
+                df.reset_index().set_index(columns, inplace=True, drop=False, **self.pk)
+            except Exception as err:  # pylint: disable=W0703
+                self._logger.exception(f"TableInput: {err}")
+        self._result = df
+        numrows = len(df.index)
+        self._variables["_numRows_"] = numrows
+        self._variables[f"{self.StepName}_NUMROWS"] = numrows
+        self.add_metric("NUM_ROWS", df.shape[0])
+        self.add_metric("NUM_COLUMNS", df.shape[1])
+        if hasattr(self, "tablename"):
+            self.add_metric("Table", self.tablename)
+            self.add_metric("Schema", self.schema)
+        else:
+            self.add_metric("Query", self.query)
+        ## result
+        self._result = df
+        return self._result