flowtask 5.8.4__cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl

flowtask/components/QueryIterator.py

@@ -0,0 +1,273 @@
+import asyncio
+from decimal import Decimal
+from typing import Any
+from collections.abc import Callable
+from pathlib import Path, PurePath
+import aiofiles
+import numpy as np
+import pandas as pd
+from navconfig.logging import logging
+from asyncdb.exceptions import NoDataFound, ProviderError
+from querysource.queries.qs import QS
+from ..conf import TASK_PATH
+from ..exceptions import ComponentError, NotSupported, DataNotFound, FileError
+from .IteratorBase import IteratorBase
+from ..interfaces import TemplateSupport, DBSupport
+
+
+class QueryIterator(DBSupport, TemplateSupport, IteratorBase):
+    """
+    QueryIterator
+
+
+    Overview
+
+        This component creates a Pandas Iterator from a QuerySource query,
+        allowing for the iteration over data returned from a query.
+
+    .. table:: Properties
+    :widths: auto
+
+
+    +------------------------+----------+-----------+--------------------------------------------------------+
+    | Name                   | Required | Summary                                                            |
+    +------------------------+----------+-----------+--------------------------------------------------------+
+    | file_sql               |   No     | The SQL file to read the query from.                               |
+    +------------------------+----------+-----------+--------------------------------------------------------+
+    | query_slug             |   No     | The query slug to use for fetching data.                           |
+    +------------------------+----------+-----------+--------------------------------------------------------+
+    | query                  |   No     | The raw SQL query string to execute.                               |
+    +------------------------+----------+-----------+--------------------------------------------------------+
+    | conditions             |   No     | Conditions to apply to the query.                                  |
+    +------------------------+----------+-----------+--------------------------------------------------------+
+    | columns                |   No     | Specific columns to extract from the result.                       |
+    +------------------------+----------+-----------+--------------------------------------------------------+
+    | use_template           |   No     | If True, use a query template for the query. Default is False.     |
+    +------------------------+----------+-----------+--------------------------------------------------------+
+    | drop_empty             |   No     | If True, drop empty rows and columns from the DataFrame.           |
+    +------------------------+----------+-----------+--------------------------------------------------------+
+    | dropna                 |   No     | Subset of columns to check for NaN values before dropping rows.    |
+    +------------------------+----------+-----------+--------------------------------------------------------+
+
+    Returns
+
+    This component returns the last value generated by the iteration process. Typically,
+    it returns a pandas DataFrame containing the query results.
+    """
+
+    def __init__(
+        self,
+        loop: asyncio.AbstractEventLoop = None,
+        job: Callable = None,
+        stat: Callable = None,
+        **kwargs,
+    ):
+        self.pk = []
+        self.data = None
+        self._iterator: Any = None
+        self._variables = {}
+        self.vars = {}
+        self._columns = []
+        self._query: str = None
+        self.use_template: bool = bool(kwargs.get("use_template", False))
+        super().__init__(loop=loop, job=job, stat=stat, **kwargs)
+
+    async def open_sqlfile(self, file: PurePath, **kwargs) -> str:
+        if file.exists() and file.is_file():
+            content = None
+            # open SQL File:
+            async with aiofiles.open(file, "r+") as afp:
+                content = await afp.read()
+                # check if we need to replace masks
+            if hasattr(self, "masks"):
+                content = self.mask_replacement(content)
+            if self.use_template is True:
+                content = self._templateparser.from_string(content, kwargs)
+            return content
+        else:
+            raise FileError(f"{self.__name__}: Missing SQL File: {file}")
+
+    async def start(self, **kwargs):
+        """Getting kind of Query."""
+        if hasattr(self, "file_sql"):
+            try:
+                file_path = Path(TASK_PATH).joinpath(
+                    self._program, "sql", self.file_sql
+                )
+                if file_path.exists() and file_path.is_file():
+                    self._query = await self.open_sqlfile(file_path)
+            except Exception as e:
+                raise FileError(
+                    f"File SQL doesn't exists: {self.file_sql!s}, {e}"
+                ) from e
+        elif hasattr(self, "query_slug"):
+            self.query_slug = self.mask_replacement(self.query_slug)
+            self._query = self.query_slug
+        elif hasattr(self, "query"):
+            self._query = self.query
+            if hasattr(self, "masks"):
+                self._query = self.mask_replacement(self._query)
+                try:
+                    self._query = self._query.format(**self._variables)
+                except Exception as err:
+                    print("Error replacing Vars in Query: ", err)
+        if hasattr(self, "conditions"):
+            self.set_conditions("conditions")
+        await super().start(**kwargs)
+        self.processing_credentials()
+        return True
+
+    async def close(self, job=None):
+        close = getattr(job, "close", None)
+        if job:
+            if asyncio.iscoroutinefunction(close):
+                await job.close()
+            else:
+                job.close()
+
+    def createJob(self, target, params, row):
+        """Create the Job Component."""
+        self._result = self.data
+        dt = {}
+        for column in self._columns:
+            value = row[column]
+            if isinstance(value, (int, np.int64, np.integer)):
+                value = int(value)
+            elif isinstance(value, (float, Decimal)):
+                value = float(value)
+            self.setVar(column, value)
+            # print('ITER ', value, type(value))
+            params[column] = value
+            dt[column] = value
+        for name, value in self.vars.items():
+            # TODO: check this logic
+            print("VARS: ", name, value, column)
+            val = row[column]
+            # print('VAL ', val)
+            # need to build this attribute
+            if isinstance(value, list):
+                pass
+                # TODO: logic to use functions with dataframes
+            #     # need to calculate the value
+            else:
+                if "{" in str(value):
+                    value = value.format(**dt)
+                else:
+                    value = val
+            params[name] = value
+            self.setVar(name, value)
+        return self.get_job(target, **params)
+
+    async def run(self):
+        """Async Run Method."""
+        # first: getting data:
+        df = None
+        result = None
+        if not self._query:
+            raise ComponentError("QueryToPandas: Empty Query/Slug or File")
+        if hasattr(self, "query") or hasattr(self, "file_sql"):
+            try:
+                connection = self.get_connection(event_loop=self._loop)
+            except Exception as err:
+                self._logger.exception(err, stack_info=True)
+                raise
+            async with await connection.connection() as conn:
+                try:
+                    res, error = await conn.query(self._query)
+                    if error:
+                        self._logger.error(f"QueryIterator: {error}")
+                        raise NoDataFound(error)
+                    result = [dict(row) for row in res]
+                    df = await self.get_dataframe(result)
+                except NoDataFound:
+                    result = []
+                except Exception as err:
+                    self._logger.error(err)
+                    raise
+        elif hasattr(self, "query_slug"):
+            conditions = {}
+            if hasattr(self, "conditions"):
+                conditions = self.conditions
+            result = await self.get_query(self._query, conditions)
+            df = await self.get_dataframe(result)
+        else:
+            raise NotSupported(f"{self.__name__}: Method not allowed")
+        # getting the iterator:
+        if not hasattr(self, "columns"):
+            # iterate over the total columns of dataframe
+            self._columns = df.columns
+        else:
+            self._columns = self.columns
+        self._iterator = df.iterrows()
+        # iterate over next task
+        step, target, params = self.get_step()
+        step_name = step.name
+        for index, row in self._iterator:
+            self._logger.debug(f"ITER: index:{index} row: {row}")
+            # iterate over every row
+            # get I got all values, create a job:
+            job = self.createJob(target, params, row)
+            # print('JOB: ', job)
+            if job:
+                try:
+                    self._result = await self.async_job(job, step_name)
+                except (NoDataFound, DataNotFound) as err:
+                    # its a data component a no data was found
+                    self._logger.debug(f"Data not Found for Task {step_name}, got: {err}")
+                    continue
+                except (ProviderError, ComponentError) as err:
+                    raise ComponentError(f"Error on {step_name}, error: {err}") from err
+                except NotSupported as err:
+                    raise NotSupported(f"Not Supported: {err}") from err
+                except Exception as err:
+                    raise ComponentError(
+                        f"Component Error {step_name}, error: {err}"
+                    ) from err
+                finally:
+                    await self.close(job)
+        # returning last value generated by iteration
+        return self._result
+
+    async def get_dataframe(self, result) -> pd.DataFrame:
+        try:
+            df = pd.DataFrame(result)
+        except Exception as err:
+            self._logger.exception(err, stack_info=True)
+        # Attempt to infer better dtypes for object columns.
+        df.infer_objects()
+        df = df.convert_dtypes(convert_string=True)
+        if hasattr(self, "drop_empty"):
+            df.dropna(axis=1, how="all", inplace=True)
+            df.dropna(axis=0, how="all", inplace=True)
+        if hasattr(self, "dropna"):
+            df.dropna(subset=self.dropna, how="all", inplace=True)
+        return df
+
+    async def get_query(self, slug, conditions: dict = None):
+        result: Any = []
+        if not conditions:
+            conditions = self.conditions
+        try:
+            qry = QS(slug=slug, conditions=conditions, loop=self._loop, lazy=True)
+            await qry.get_query()
+        except (DataNotFound, NoDataFound) as ex:
+            raise DataNotFound(f"{ex!s}") from ex
+        except Exception as err:
+            raise ComponentError(f"{err}") from err
+        try:
+            res, error = await qry.query()
+            if not res:
+                raise NoDataFound(f"{slug}: Data Not Found")
+            if error:
+                raise ComponentError(f"Error on Query: {error}")
+            result = result + [dict(row) for row in res]
+            return result
+        except (DataNotFound, NoDataFound) as err:
+            raise DataNotFound(f"{err!s}") from err
+        except Exception as err:
+            raise ComponentError(f"Error on Query: {err}") from err
+        finally:
+            try:
+                await qry.close()
+            except Exception as err:
+                self._logger.exception(err, stack_info=True)

flowtask/components/QueryToInsert.py

@@ -0,0 +1,327 @@
+# -*- coding: utf-8 -*-
+from typing import Any
+from collections.abc import Callable
+import asyncio
+from datetime import datetime
+from pathlib import PosixPath
+from tqdm import tqdm
+import orjson
+from datamodel.parsers.json import json_encoder
+from asyncdb.drivers.postgres import postgres
+from asyncdb.exceptions import ProviderError
+from ..exceptions import ComponentError
+from ..utils import SafeDict
+from ..conf import default_dsn
+from .flow import FlowComponent
+
+
+class QueryToInsert(FlowComponent):
+    """
+    QueryToInsert.
+
+
+    Overview
+
+        This component allows me to insert data into a database schema
+
+    .. table:: Properties
+       :widths: auto
+
+
+    +--------------+----------+-----------+-------------------------------------------------------+
+    | Name         | Required | Summary                                                           |
+    +--------------+----------+-----------+-------------------------------------------------------+
+    | schema       |   Yes    | Name of the schema where is to the table                          |
+    +--------------+----------+-----------+-------------------------------------------------------+
+    | tablename    |   Yes    | Name of the table in the database                                 |
+    +--------------+----------+-----------+-------------------------------------------------------+
+    | action       |   Yes    | Sets the action to execute in this case an insert                 |
+    +--------------+----------+-----------+-------------------------------------------------------+
+    | pk           |   Yes    | Primary key to the table in the database                          |
+    +--------------+----------+-----------+-------------------------------------------------------+
+    | directory    |   Yes    | Source directory where the file is located                        |
+    +--------------+----------+-----------+-------------------------------------------------------+
+    | filter       |   Yes    | This attribute allows me to apply a filter to the data            |
+    +--------------+----------+-----------+-------------------------------------------------------+
+
+
+
+    Return the list of arbitrary days
+
+    
+
+        Example:
+
+        ```yaml
+        QueryToInsert:
+          schema: public
+          tablename: queries
+          action: insert
+          pk:
+          - query_slug
+          directory: /home/ubuntu/symbits/
+          filter:
+            query_slug: walmart_stores
+        ```
+
+    """
+    """
+    QueryToInsert
+
+    Overview
+
+        This component generates SQL INSERT or UPDATE statements from a query and saves them to a file.
+
+    .. table:: Properties
+    :widths: auto
+
+
+    +------------------------+----------+-----------+-------------------------------------------------------+
+    | Name                   | Required | Summary                                                           |
+    +------------------------+----------+-----------+-------------------------------------------------------+
+    | schema                 |   Yes    | The schema of the table to insert or update data in.              |
+    +------------------------+----------+-----------+-------------------------------------------------------+
+    | tablename              |   Yes    | The name of the table to insert or update data in.                |
+    +------------------------+----------+-----------+-------------------------------------------------------+
+    | action                 |   Yes    | The action to perform, either "insert" or "update".               |
+    +------------------------+----------+-----------+-------------------------------------------------------+
+    | pk                     |   Yes    | The primary key(s) of the table.                                  |
+    +------------------------+----------+-----------+-------------------------------------------------------+
+    | filter                 |   No     | Filters to apply to the query.                                    |
+    +------------------------+----------+-----------+-------------------------------------------------------+
+    | directory              |   Yes    | The directory to save the SQL file.                               |
+    +------------------------+----------+-----------+-------------------------------------------------------+
+    | fields                 |   No     | Specific fields to select in the query.                           |
+    +------------------------+----------+-----------+-------------------------------------------------------+
+
+    Returns
+
+    This component returns True if the SQL file is successfully created, otherwise raises a ComponentError.
+    """
+    where = ""
+    query = "SELECT {fields} FROM {schema}.{table} {where} order by ctid;"
+    insert = "INSERT INTO {}.{} ({}) VALUES({}) {};\n"
+    conflict = "ON CONFLICT ({}) DO UPDATE SET {}"
+    conflict_update = "{} = EXCLUDED.{}"
+    update = "UPDATE {}.{} SET {} {};\n"
+
+    def __init__(
+        self,
+        loop: asyncio.AbstractEventLoop = None,
+        job: Callable = None,
+        stat: Callable = None,
+        **kwargs,
+    ):
+        """Init Method."""
+        self.action: str = None
+        self.schema: str = ""
+        self.tablename: str = ""
+        self.pk: Any = None
+        self.filter: Any = None
+        self._connection: Callable = None
+        self._fields = kwargs.pop('fields', [])
+        super(QueryToInsert, self).__init__(
+            loop=loop,
+            job=job,
+            stat=stat,
+            **kwargs
+        )
+
+    def _escapeString(self, value):
+        v = value if value != "None" or value is not None else ""
+        v = str(v).replace("'", "''") if type(v) == str else v
+        v = "'{}'".format(v) if type(v) == str else v
+        v = "array{}".format(v) if type(v) == list else v
+        return v
+
+    async def start(self, **kwargs):
+        """Start."""
+        if not hasattr(self, "schema"):
+            raise ComponentError("Schema not defined:")
+        if not hasattr(self, "tablename"):
+            raise ComponentError("Tablename not defined:")
+        if not hasattr(self, "action"):
+            raise ComponentError("Action not defined:")
+        if not hasattr(self, "pk"):
+            raise ComponentError("Primary keys not defined:")
+        else:
+            if isinstance(self.pk, str):
+                try:
+                    self.pk = orjson.loads(self.pk)
+                except Exception:
+                    self.pk = self.pk
+        if not hasattr(self, "directory"):
+            raise ComponentError("Directory not defined:")
+        if hasattr(self, "filter"):
+            filter = []
+            if isinstance(self.filter, str):
+                self.filter = orjson.loads(self.filter)
+            for key, value in self.filter.items():
+                filter.append(
+                    "{} = {}".format(
+                        key, ("'{}'".format(value) if type(value) == str else value)
+                    )
+                )
+            if len(filter) > 0:
+                self.where = "WHERE {}".format(" AND ".join(filter))
+        today = datetime.today().strftime("%Y-%m-%d")
+        # Create directory if not exists
+        try:
+            PosixPath(self.directory).mkdir(parents=True, exist_ok=True)
+        except Exception as err:
+            self.logger.error(f"Error creating directory {self.directory}: {err}")
+            raise ComponentError(
+                f"Error creating directory {self.directory}: {err}"
+            ) from err
+        self.filename = PosixPath(
+            self.directory,
+            "{}_{}_{}.{}.sql".format(today, self.action, self.schema, self.tablename),
+        )
+        if not self._fields:
+            fields = '*'
+        else:
+            fields = ', '.join(self._fields)
+        self.query = self.query.format_map(
+            SafeDict(fields=fields)
+        )
+        self.query = self.query.format(
+            schema=self.schema,
+            table=self.tablename,
+            where=self.where
+        )
+        self._logger.notice(
+            f"Query: {self.query}"
+        )
+        return True
+
+    async def get_connection(self):
+        try:
+            self._connection = postgres(dsn=default_dsn, loop=self._loop)
+            await self._connection.connection()
+        except Exception as err:
+            raise ProviderError(f"Error configuring pg Connection: {err!s}") from err
+        return self._connection
+
+    async def run(self):
+        """Async Method."""
+        # get connection
+        self._connection = await self.get_connection()
+        res, err = await self._connection.query(self.query)
+        if err:
+            raise ComponentError(
+                'Query Error "{}": {}'.format(self.query, err)
+            )
+        colinfo = self.column_info("{}.{}".format(self.schema, self.tablename))
+        if res:
+            query = []
+            if self.action == "insert":
+                total = len(res)
+                with tqdm(total=total) as pbar:
+                    for row in res:
+                        columns_update = row.keys()
+                        columns = row.keys()
+                        values = []
+                        for col in row.keys():
+                            values.append(self.get_values(col, colinfo[col], row[col]))
+                        update = ", ".join(
+                            [self.conflict_update.format(c, c) for c in columns_update]
+                        )
+                        conflict = self.conflict.format(",\n".join(self.pk), update)
+                        query.append(
+                            self.insert.format(
+                                self.schema,
+                                self.tablename,
+                                ", ".join(columns),
+                                ", ".join(values),
+                                conflict,
+                            )
+                        )
+                        pbar.update(1)
+            elif self.action == "update":
+                for row in res:
+                    columns = row.keys()
+                    values = []
+                    for col in row.keys():
+                        values.append(self.get_values(col, colinfo[col], row[col]))
+                    vals = dict(zip(list(columns), values))
+                    pkwhere = "AND".join(["{}={}".format(v, vals[v]) for v in self.pk])
+                    where = (
+                        "{} AND {}".format(self.where, pkwhere)
+                        if self.where != ""
+                        else "WHERE {}".format(pkwhere)
+                    )
+                    update_list = ["{}={}".format(k, v) for k, v in vals.items()]
+                    query.append(
+                        self.update.format(
+                            self.schema, self.tablename, ", ".join(update_list), where
+                        )
+                    )
+
+            f = open(self.filename, "w")
+            f.write("".join(query))
+            f.close()
+            self._result = "".join(query)
+            self.add_metric('NUM_ROWS', len(query))
+            self.add_metric('QUERY', self.query)
+            # return self._result
+            return True
+        else:
+            raise ComponentError(
+                "Error creating Query File: Empty Result."
+            )
+
+    def close(self):
+        """Method."""
+        pass
+
+    def column_info(self, table):
+        result = None
+        sql = f"""
+        SELECT attname AS column_name, atttypid::regtype AS data_type, attnotnull::boolean as notnull
+        FROM pg_attribute WHERE attrelid = '{table}'::regclass AND attnum > 0 AND NOT attisdropped
+        ORDER  BY attnum"""
+        db = postgres(dsn=default_dsn)
+        conn = db.connect()
+        result, error = conn.fetchall(sql)
+        if error:
+            raise ComponentError(f"Error executing query: {error}")
+        if result:
+            return {item["column_name"]: item["data_type"] for item in result}
+        else:
+            print(f"The table {self.schema}.{self.tablename} does not exist")
+            return None
+
+    def get_values(self, column, type, value):
+        if value is None or value == "None":
+            return "NULL"
+        elif type in [
+            "text",
+            "character varying",
+            "timestamp with time zone",
+            "timestamp without time zone",
+            "date",
+            "uuid",
+            "inet",
+        ]:
+            return "'{}'".format(str(value).replace("'", "''"))
+        elif type in ["hstore"]:
+            return "'{}'".format(
+                ", ".join(
+                    [
+                        '"{}"=>"{}"'.format(k, v).replace("'", "''")
+                        for k, v in value.items()
+                    ]
+                )
+            )
+        elif type in ["jsonb"]:
+            return "'{}'".format(json_encoder(value).replace("'", "''"))
+        elif type in ["character varying[]"]:
+            return "'{{{}}}'".format(
+                ", ".join(['"{}"'.format(v) for v in value]).replace("'", "''")
+            )
+        elif type in ["integer[]"]:
+            return "'{{{}}}'".format(
+                ", ".join([str(v) for v in value]).replace("'", "''")
+            )
+        else:
+            return str(value)