flowtask 5.8.4__cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl

flowtask/components/RethinkDBQuery.py

@@ -0,0 +1,189 @@
+
+import asyncio
+from typing import List, Dict, Union
+from collections.abc import Callable
+import rethinkdb as r
+from ..exceptions import ComponentError, DataNotFound, ConfigError
+from .flow import FlowComponent
+from ..interfaces.dataframes import PandasDataframe
+from ..interfaces import TemplateSupport
+from ..interfaces.databases.rethink import RethinkDBSupport
+
+
+class RethinkDBQuery(
+    RethinkDBSupport,
+    FlowComponent,
+    PandasDataframe,
+    TemplateSupport,
+):
+    """
+    RethinkDBQuery.
+
+    Class to execute queries against a RethinkDB database and retrieve results.
+    using asyncDB as backend.
+
+    RethinkDB Query can support queries by mapping RethinkDB methods as attributes.
+    Methods as "table", "filter", "order_by", "limit", "pluck" are supported.
+
+    
+
+        Example:
+
+        ```yaml
+        RethinkDBQuery:
+          table: stores_reviews
+          schema: epson
+          filter:
+          - rating:
+              gt: 4
+          - rating:
+              lt: 6
+          order_by:
+          - rating: desc
+          limit: 50
+          columns:
+          - store_id
+          - store_name
+          - formatted_address
+          - latitude
+          - longitude
+          - reviews
+          - rating
+          - user_ratings_total
+          as_dataframe: true
+        ```
+
+    """
+    def __init__(
+        self,
+        loop: asyncio.AbstractEventLoop = None,
+        job: Callable = None,
+        stat: Callable = None,
+        **kwargs,
+    ) -> None:
+        """Init Method."""
+        self.table = kwargs.get('table', None)
+        self.schema = kwargs.get('schema', None)
+        super().__init__(loop=loop, job=job, stat=stat, **kwargs)
+        self._db = None
+
+    async def close(self):
+        """Close the connection to the RethinkDB database."""
+        if self._db:
+            try:
+                await self._db.close()
+            except Exception:
+                pass
+        self._db = None
+
+    async def start(self, **kwargs):
+        await super().start(**kwargs)
+        if not hasattr(self, 'table'):
+            raise ConfigError(
+                "'table' attribute is required in the RethinkDBQuery."
+            )
+        if not hasattr(self, 'schema'):
+            raise ConfigError(
+                "'schema' attribute is required in the RethinkDBQuery."
+            )
+        # Replacing with Masking if needed.
+        self.schema = self.mask_replacement(self.schema)
+        self.table = self.mask_replacement(self.table)
+        if hasattr(self, 'columns'):
+            # used as "pluck"
+            self.pluck = self.columns
+        return True
+
+    def _filter_criteria(self, engine, cursor):
+        result = engine.expr(True)
+        for args in self.filter:
+            field, inner_args = next(iter(args.items()))
+            func, value = next(iter(inner_args.items()))
+            if func == 'in':
+                cursor = cursor.filter(
+                    (
+                        lambda exp: engine.expr(value)
+                        .coerce_to("array")
+                        .contains(exp[field])
+                    )
+                )
+            elif func == 'gt':
+                result = result.and_(
+                    engine.row[field].gt(value)
+                )
+            elif func == 'eq':
+                result = result.and_(
+                    engine.row[field].eq(value)
+                )
+            elif func == 'lt':
+                result = result.and_(
+                    engine.row[field].lt(value)
+                )
+            elif func == 'ge':
+                result = result.and_(
+                    engine.row[field].ge(value)
+                )
+            elif func == 'le':
+                result = result.and_(
+                    engine.row[field].le(value)
+                )
+        cursor = cursor.filter(result)
+        return cursor
+
+    def _order_by(self, engine, cursor):
+        order_clauses = []
+        for order in self.order_by:
+            field, direction = next(iter(order.items()))
+            if direction.lower() == 'asc':
+                order_clauses.append(engine.asc(field))
+            elif direction.lower() == 'desc':
+                order_clauses.append(engine.desc(field))
+            else:
+                raise ComponentError(
+                    f"Invalid order direction: {direction}"
+                )
+        if order_clauses:
+            cursor = cursor.order_by(*order_clauses)
+        return cursor
+
+    async def run(self):
+        """Execute the RethinkDB query and retrieve the results."""
+        if not self._db:
+            # TODO: add support for datasources.
+            self._db = self.default_connection()
+        try:
+            async with await self._db.connection() as conn:
+                # Change to default database:
+                engine = conn.engine()
+                # changing to active database
+                cursor = engine.db(self.schema).table(self.table)
+                if hasattr(self, 'filter'):
+                    # Build a Filter functionality:
+                    cursor = self._filter_criteria(engine, cursor)
+                if hasattr(self, 'order_by'):
+                    cursor = self._order_by(engine, cursor)
+                    # cursor = cursor.order_by(self.order_by)
+                    pass
+                if hasattr(self, 'limit'):
+                    cursor = cursor.limit(self.limit)
+                if hasattr(self, 'pluck'):
+                    cursor = cursor.pluck(self.pluck)
+                data = []
+                print('CURSOR > ', cursor)
+                cursor = await cursor.run(conn.get_connection())
+                if isinstance(cursor, list):
+                    data = cursor
+                else:
+                    while await cursor.fetch_next():
+                        item = await cursor.next()
+                        data.append(item)
+            # Check if return as Dataframe:
+            if self.as_dataframe is True:
+                self._result = await self.create_dataframe(data)
+            else:
+                self._result = data
+            return self._result
+        except Exception as e:
+            raise ComponentError(
+                f"Error executing RethinkDB query: {e}"
+            ) from e

flowtask/components/Rsync.py

@@ -0,0 +1,74 @@
+from collections.abc import Callable
+import asyncio
+import asyncssh
+from ..exceptions import ComponentError
+from .RunSSH import RunSSH
+
+
+class Rsync(RunSSH):
+    def __init__(
+        self,
+        loop: asyncio.AbstractEventLoop = None,
+        job: Callable = None,
+        stat: Callable = None,
+        **kwargs,
+    ):
+        self.flags: str = "azrv"
+        super(Rsync, self).__init__(loop=loop, job=job, stat=stat, **kwargs)
+
+    async def start(self, **kwargs):
+        super(Rsync, self).start(**kwargs)
+        if hasattr(self, "source"):  # using the destination filosophy
+            try:
+                if hasattr(self, "masks"):
+                    self.source_dir = self.mask_replacement(self.source["directory"])
+                else:
+                    self.source_dir = self.source["directory"]
+            except KeyError as exc:
+                raise ComponentError(
+                    "Rsync Error: you must specify a source directory"
+                ) from exc
+        if hasattr(self, "destination"):
+            if hasattr(self, "masks"):
+                self.destination_dir = self.mask_replacement(
+                    self.destination["directory"]
+                )
+            else:
+                self.destination_dir = self.destination["directory"]
+            # also, calculate the destination server:
+            self.dest_server = self.mask_replacement(self.destination["server"])
+            self.dest_user = self.mask_replacement(self.destination["user"])
+            try:
+                self.dest_port = self.mask_replacement(self.destination["port"])
+            except KeyError:
+                self.dest_port = None
+        return True
+
+    async def run(self):
+        await self.open(
+            host=self.host,
+            port=self.port,
+            tunnel=self.tunnel,
+            credentials=self.credentials,
+        )
+        rsync = "rsync -{flags} {source} {destination}"
+        if self.dest_port is not None:
+            rsync = rsync + f" --port={self.dest_port}"
+        destination = f"{self.dest_user}@{self.dest_server}:{self.destination_dir}"
+        command = rsync.format(
+            flags=self.flags, destination=destination, source=self.source_dir
+        )
+        try:
+            rst = await self._connection.run(command, check=True)
+            result = {
+                "exit_status": rst.exit_status,
+                "returncode": rst.returncode,
+                "error": rst.stderr,
+                # "stdout": rst.stdout
+            }
+        except asyncssh.process.ProcessError as err:
+            self._logger.error(f"Error executing command: {err}")
+        self.add_metric("SSH: COMMAND", command)
+        self.add_metric("SSH: RESULT", result)
+        self._result = result
+        return result

flowtask/components/RunSSH.py

@@ -0,0 +1,59 @@
+import asyncio
+from collections.abc import Callable
+import logging
+import asyncssh
+from .flow import FlowComponent
+from ..interfaces.SSHClient import SSHClient
+
+
+class RunSSH(SSHClient, FlowComponent):
+    """
+    RunSSH.
+
+    Run any arbitrary command into an SSH server.
+    """
+
+    def __init__(
+        self,
+        loop: asyncio.AbstractEventLoop = None,
+        job: Callable = None,
+        stat: Callable = None,
+        **kwargs,
+    ):
+        super().__init__(loop=loop, job=job, stat=stat, **kwargs)
+
+    async def start(self, **kwargs):
+        """Start.
+
+        Processing variables and credentials.
+        """
+        try:
+            self.define_host()
+            self.processing_credentials()
+        except Exception as err:
+            logging.error(err)
+            raise
+
+    async def run(self):
+        result = {}
+        await self.open(
+            host=self.host,
+            port=self.port,
+            tunnel=self.tunnel,
+            credentials=self.credentials,
+        )
+        for command in self.commands:
+            command = self.mask_replacement(command)
+            try:
+                rst = await self._connection.run(command, check=True)
+                result[command] = {
+                    "exit_status": rst.exit_status,
+                    "returncode": rst.returncode,
+                    "error": rst.stderr,
+                    # "stdout": rst.stdout
+                }
+            except asyncssh.process.ProcessError as err:
+                logging.error(f"Error executing command: {err}")
+        self.add_metric("SSH: COMMAND", result)
+        self._result = result
+        return result

flowtask/components/RunShell.py

@@ -0,0 +1,71 @@
+import subprocess
+import asyncio
+from typing import List
+from collections.abc import Callable
+from navconfig.logging import logging
+from .flow import FlowComponent
+
+
+class RunShell(FlowComponent):
+    """
+    RunShell.
+
+
+    Overview
+
+        Execute a Command to run a task
+
+    .. table:: Properties
+    :widths: auto
+
+
+    +--------------+----------+-----------+-------------------------------------------------------+
+    | Name         | Required | Summary                                                           |
+    +--------------+----------+-----------+-------------------------------------------------------+
+    | name         |   Yes    | Name of task                                                      |
+    +--------------+----------+-----------+-------------------------------------------------------+
+    | description  |   Yes    | Task description                                                  |
+    +--------------+----------+-----------+-------------------------------------------------------+
+    | steps        |   Yes    | Not assigned steps                                                |
+    +--------------+----------+-----------+-------------------------------------------------------+
+    | runtask      |   Yes    | This method runs the task                                         |
+    +--------------+----------+-----------+-------------------------------------------------------+
+    | program      |   Yes    | Program name                                                      |
+    +--------------+----------+-----------+-------------------------------------------------------+
+    | task         |   Yes    | Assign the run shell attribute                                    |
+    +--------------+----------+-----------+-------------------------------------------------------+
+
+    Return the list of arbitrary days
+    """
+
+    def __init__(
+        self,
+        loop: asyncio.AbstractEventLoop = None,
+        job: Callable = None,
+        stat: Callable = None,
+        **kwargs,
+    ):
+        """Init Method."""
+        self.commands: List = []
+        super(RunShell, self).__init__(loop=loop, job=job, stat=stat, **kwargs)
+
+    async def start(self, **kwargs):
+        return True
+
+    async def close(self):
+        pass
+
+    async def run(self):
+        for cmd in self.commands:
+            if hasattr(self, "masks"):
+                cmd = self.mask_replacement(cmd)
+                logging.debug(">", cmd)
+            try:
+                result = subprocess.Popen(
+                    cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE
+                ).communicate()
+                logging.debug(result)
+                return True
+            except subprocess.CalledProcessError as e:
+                print(f"Error in command: {e}")
+                return False

flowtask/components/SalesForce.py

@@ -0,0 +1,20 @@
+from navconfig.logging import logging
+from querysource.exceptions import DataNotFound as QSNotFound
+from ..exceptions import ComponentError, DataNotFound
+from .QSBase import QSBase
+
+
+class SalesForce(QSBase):
+    """SalesForce Connector."""
+
+    type = "report"
+    _driver = "salesforce"
+
+    async def report(self):
+        try:
+            return await self._qs.report()
+        except QSNotFound as err:
+            raise DataNotFound(f"SalesForce: Report Not Found: {err}") from err
+        except Exception as err:
+            logging.exception(err)
+            raise ComponentError(f"SalesForce ERROR: {err!s}") from err

flowtask/components/SaveImageBank/__init__.py

@@ -0,0 +1,257 @@
+from typing import Optional
+from collections.abc import Callable
+import re
+import asyncio
+import asyncpg
+from PIL.TiffImagePlugin import IFDRational
+from pgvector.asyncpg import register_vector
+from querysource.types.validators import Entity
+from navigator.libs.json import JSONContent
+from ..flow import FlowComponent
+from ...exceptions import ConfigError, ComponentError
+from ...conf import default_dsn
+
+IDENT_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
+
+def qid(name: str) -> str:
+    """
+    Very small helper to quote SQL identifiers safely.
+    Raises if name contains anything but letters, digits or '_'.
+    """
+    if not IDENT_RE.match(name):
+        raise ValueError(
+            f"illegal identifier: {name!r}"
+        )
+    return '"' + name + '"'
+
+class SaveImageBank(FlowComponent):
+    """
+    SaveImageBank.
+
+    Save images into a postgreSQL Table, with UPSERT and optional evaluation for duplicates.
+    """
+    def __init__(
+        self,
+        loop: Optional[asyncio.AbstractEventLoop] = None,
+        job: Callable | None = None,
+        stat: Callable | None = None,
+        **kwargs,
+    ):
+        self.id_column: str = kwargs.get("id_column", "photo_id")
+        self.hash_column: str = kwargs.get("hash_column", "image_hash")
+        self.vector_column: str = kwargs.get("vector_column", "image_vector")
+        self.detections_column: str = kwargs.get("detections_column", "image_features")
+        self.hamming_threshold: int = kwargs.get("hamming_threshold", 4)
+        self.vector_threshold: float = kwargs.get("vector_threshold", 0.05)
+        self.tablename: str = kwargs.get("tablename", "image_bank")
+        self.schema: str = kwargs.get("schema", "public")
+        self.pool: asyncpg.Pool | None = None
+        self._semaphore = asyncio.Semaphore(16)  # limit GPU tasks
+        self.drop_columns: list[str] = kwargs.get("drop_columns", [])
+        super().__init__(loop=loop, job=job, stat=stat, **kwargs)
+        # JSON encoder:
+        self._encoder = JSONContent()
+
+    def _qualified_tablename(self) -> str:
+        """
+        Get the qualified table name.
+        """
+        if not self.schema:
+            raise ConfigError("Schema is not set.")
+        if not self.tablename:
+            raise ConfigError("Table name is not set.")
+        return f"{qid(self.schema)}.{qid(self.tablename)}"
+
+    def _build_insert_sql(self, columns: list[str]) -> str:
+        """
+        Produces something like:
+
+        INSERT INTO schema.table (col1,col2,…) VALUES ($1,$2,…)
+        ON CONFLICT (photo_id) DO UPDATE SET
+            col1 = EXCLUDED.col1,
+            ...
+        """
+        col_list = ", ".join(map(qid, columns))
+        placeholders = ", ".join(f"${i}" for i in range(1, len(columns) + 1))
+        updates = ", ".join(f"{qid(c)} = EXCLUDED.{qid(c)}" for c in columns
+                            if c != self.id_column)
+
+        return (
+            f"INSERT INTO {self._qualified_tablename()} ({col_list}) "
+            f"VALUES ({placeholders}) "
+            f"ON CONFLICT ({qid(self.id_column)}) "
+            f"DO UPDATE SET {updates};"
+        )
+
+    def _build_phash_sql(self) -> str:
+        return (
+            f"SELECT 1 FROM {self._qualified_tablename()} "
+            f"WHERE {qid(self.id_column)} IS DISTINCT FROM $3 "
+            f"AND   bit_count(('x' || $1)::bit(256) # "
+            f"                 ('x' || {qid(self.hash_column)})::bit(256)) "
+            f"      <= $2 "
+            f"LIMIT 1;"
+        )
+
+    def _build_vector_sql(self) -> str:
+        return (
+            f"SELECT 1 FROM {self._qualified_tablename()} "
+            f"WHERE {qid(self.id_column)} IS DISTINCT FROM $3 "
+            f"AND   {qid(self.vector_column)} <#> $1::vector < $2 "
+            f"LIMIT 1;"
+        )
+
+    async def pgvector_init(self, conn):
+        """
+        Initialize pgvector extension in PostgreSQL.
+        """
+        # Setup jsonb encoder/decoder
+        def _encoder(value):
+            # return json.dumps(value, cls=BaseEncoder)
+            return self._encoder.dumps(value)  # pylint: disable=E1120
+
+        def _decoder(value):
+            return self._encoder.loads(value)  # pylint: disable=E1120
+
+        await conn.set_type_codec(
+            "json",
+            encoder=_encoder,
+            decoder=_decoder,
+            schema="pg_catalog"
+        )
+        await conn.set_type_codec(
+            "jsonb",
+            encoder=_encoder,
+            decoder=_decoder,
+            schema="pg_catalog"
+        )
+
+        await register_vector(conn)
+
+    # ──────────────────────────────────────────────────────────────
+    # Setup / teardown
+    # ──────────────────────────────────────────────────────────────
+    async def start(self, **kwargs):
+        if self.previous:
+            self.data = self.input
+
+        # column checks
+        for col in (self.id_column, self.hash_column,
+                    self.vector_column, self.detections_column):
+            if col not in self.data.columns:
+                raise ConfigError(
+                    f"Column '{col}' missing from DataFrame"
+                )
+        self.pool = await asyncpg.create_pool(
+            dsn=default_dsn,
+            min_size=1,
+            max_size=4,
+            max_queries=100,
+            init=self.pgvector_init,
+            timeout=10,
+        )
+        # Check if the table exists
+        if not self.pool:
+            raise ConfigError(
+                "Database connection pool is not initialized."
+            )
+        async with self.pool.acquire() as conn:
+            try:
+                await conn.execute(
+                    f"SELECT 1 FROM {self.schema}.{self.tablename} LIMIT 1"
+                )
+            except asyncpg.exceptions.UndefinedTableError:
+                raise ConfigError(
+                    f"Table {self.schema}.{self.tablename} does not exist."
+                )
+            except asyncpg.exceptions.UndefinedSchemaError:
+                raise ConfigError(
+                    f"Schema {self.schema} does not exist."
+                )
+        if "duplicated" not in self.data.columns:
+            self.data["duplicated"] = False
+        # prepare SQL strings
+        self._sql_phash = self._build_phash_sql()
+        self._sql_vector = self._build_vector_sql()
+
+    async def close(self):
+        if self.pool:
+            await self.pool.close()
+
+    # --------------- duplicate test --------------------
+    async def _is_duplicated(self, conn, phash: str, vec: list[float], current_id: int) -> bool:
+        """
+        Check if the given hash and vector are duplicated in the database.
+        :param conn: Database connection.
+        :param phash: Perceptual hash of the image.
+        :param vec: Vector representation of the image.
+        :return: True if the image is duplicated, False otherwise.
+        """
+        # phash first
+        phash_dup = False
+        vector_dup = False
+        if phash:
+            if await conn.fetchval(self._sql_phash, phash, self.hamming_threshold, current_id):
+                phash_dup = True
+        # vector second
+        vector_dup = bool(
+            await conn.fetchval(self._sql_vector, vec, self.vector_threshold, current_id)
+        )
+        # return True if both are duplicated
+        return phash_dup and vector_dup
+
+    async def _upsert_row(self, conn, row) -> bool:
+        """
+        UPSERT a single row into the database.
+        :param conn: Database connection.
+        :param row: Row data to be inserted/updated.
+        :return: True if the row was duplicated, False otherwise.
+        """
+        # --------------- UPSERT one row --------------------
+        phash = row[self.hash_column]
+        vec = row[self.vector_column]
+        dup = await self._is_duplicated(
+            conn,
+            phash,
+            vec,
+            current_id=row[self.id_column]
+        )
+
+        # Add/overwrite duplicated flag in the in‑memory DF row
+        row["duplicated"] = dup
+
+        # Build VALUES array in the same order as self.data.columns
+        values = [row[col] for col in self.data.columns]
+        # asyncpg needs list/tuple for pgvector, ensure np → list
+        idx_vec = self.data.columns.get_loc(self.vector_column)
+        values[idx_vec] = list(values[idx_vec])
+
+        await conn.execute(self._sql_insert, *values)
+
+    async def run(self):
+        """
+        Run the task.
+        """
+        if self.pool is None:
+            raise ConfigError("Database connection pool is not initialized.")
+        if self.drop_columns:
+            # drop columns from dataframe:
+            self.data.drop(
+                columns=self.drop_columns,
+                axis=1,
+                inplace=True,
+            )
+        #
+        self._sql_insert = self._build_insert_sql(list(self.data.columns))
+
+        # check for duplicates
+        async def handle(idx):
+            async with self._semaphore, self.pool.acquire() as conn:
+                row = self.data.loc[idx].to_dict()
+                await self._upsert_row(conn, row)
+                # write duplicated flag back into DF
+                self.data.at[idx, "duplicated"] = row["duplicated"]
+        await asyncio.gather(*(handle(i) for i in self.data.index))
+
+        self._result = self.data
+        return self._result