flowtask 5.8.4__cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl

flowtask/components/CopyTo.py

@@ -0,0 +1,177 @@
+import os
+from collections.abc import Callable
+from abc import abstractmethod
+import asyncio
+from decimal import Decimal
+import numpy as np
+import pandas as pd
+import datetime
+from ..utils import SafeDict
+from .flow import FlowComponent
+from ..exceptions import (
+    ComponentError,
+    DataNotFound,
+)
+from ..interfaces.qs import QSSupport
+
+
+dtypes = {
+    "varchar": str,
+    "character varying": str,
+    "string": str,
+    "object": str,
+    "int": int,
+    "int4": int,
+    "integer": int,
+    "bigint": np.int64,
+    "int64": np.int64,
+    "uint64": np.int64,
+    "Int8": int,
+    "float64": Decimal,
+    "float": Decimal,
+    "boolean": bool,
+    "bool": bool,
+    "datetime64[ns]": datetime.datetime,
+    "date": datetime.date,
+}
+
+class CopyTo(QSSupport, FlowComponent):
+    """CopyTo.
+
+    Abstract Class for Copying (saving) a Pandas Dataframe onto a Resource
+    (example: Copy to PostgreSQL).
+    """
+    def __init__(
+        self,
+        loop: asyncio.AbstractEventLoop = None,
+        job: Callable = None,
+        stat: Callable = None,
+        **kwargs,
+    ):
+        self.pk = []
+        self.truncate: bool = False
+        self.data = None
+        self._engine = None
+        self.tablename: str = ""
+        self.schema: str = ""
+        self.use_chunks = False
+        self.chunksize = None
+        self._connection: Callable = None
+        self._driver: str = kwargs.pop('driver', 'pg')
+        self.multi = bool(kwargs.pop('multi', False))
+        self._naive_columns: list = kwargs.pop('naive_tz', [])
+        self._json_columns: list = kwargs.pop('json_columns', [])
+        self._vector_columns: list = kwargs.pop('vector_columns', [])
+        self._array_columns: list = kwargs.pop('array_columns', [])
+        self._binary_columns: list = kwargs.pop('binary_columns', [])
+        self._atomic: bool = kwargs.pop('atomic', False)
+        super(CopyTo, self).__init__(loop=loop, job=job, stat=stat, **kwargs)
+
+    async def start(self, **kwargs):
+        """Obtain Pandas Dataframe."""
+        if hasattr(self, 'credentials'):
+            self.processing_credentials()
+        if self.previous:
+            self.data = self.input
+        else:
+            raise ComponentError(
+                "CopyTo: Data Was Not Found"
+            )
+        for attr, value in self.__dict__.items():
+            if isinstance(value, str):
+                val = value.format_map(SafeDict(**self._variables))
+                object.__setattr__(self, attr, val)
+        if not self.schema:
+            try:
+                self.schema = self._program
+            except (ValueError, AttributeError, TypeError) as ex:
+                raise ComponentError(
+                    "CopyTo: Schema name not defined."
+                ) from ex
+        # Getting the connection, DSN or credentials:
+        self._connection = await self.create_connection(
+            driver=self._driver
+        )
+
+    async def close(self):
+        """Close the connection to Database."""
+        try:
+            if self._connection:
+                await self._connection.close()
+                self._logger.info(
+                    "CopyTo: Connection closed."
+                )
+        except Exception as err:
+            self._logger.error(
+                f"CopyTo: Error closing connection: {err}"
+            )
+
+    @abstractmethod
+    async def _create_table(self):
+        pass
+
+    @abstractmethod
+    async def _truncate_table(self):
+        pass
+
+    @abstractmethod
+    async def _copy_dataframe(self):
+        pass
+
+    @abstractmethod
+    async def _copy_iterable(self):
+        pass
+
+    async def run(self):
+        """Run Copy into table functionality."""
+        self._result = None
+        if self.data is None:
+            raise DataNotFound(
+                f"{self.__name__} Error: No data in Dataframe"
+            )
+        if isinstance(self.data, pd.DataFrame) and self.data.empty:
+            raise DataNotFound(
+                f"{self.__name__} Error: No data in Dataframe"
+            )
+        # Pass through data to next component.
+        self._result = self.data
+        if isinstance(self.data, pd.DataFrame):
+            columns = list(self.data.columns)
+            self.add_metric("NUM_ROWS", self.data.shape[0])
+            self.add_metric("NUM_COLUMNS", self.data.shape[1])
+            if self._debug:
+                print(f"Debugging: {self.__name__} ===")
+                for column in columns:
+                    t = self.data[column].dtype
+                    print(column, "->", t, "->", self.data[column].iloc[0])
+        elif isinstance(self.data, list):
+            columns = list(self.data[0].keys())
+            self.add_metric("NUM_ROWS", len(self.data))
+            self.add_metric("NUM_COLUMNS", len(columns))
+        else:
+            self.add_metric("DATA_SIZE", len(self.data))
+        if hasattr(self, "create_table"):
+            # Create a Table using Model
+            self._logger.debug(
+                f":: Creating table: {self.schema}.{self.tablename}"
+            )
+            await self._create_table()
+        if self.truncate is True:
+            if self._debug:
+                self._logger.debug(
+                    f"Truncating table: {self.schema}.{self.tablename}"
+                )
+            await self._truncate_table()
+        if isinstance(self.data, pd.DataFrame):
+            # insert data directly into table
+            columns = list(self.data.columns)
+            await self._copy_dataframe()
+        else:
+            # insert data using iterable
+            await self._copy_iterable()
+        self._logger.debug(
+            f"{self.__name__}: Saving results into: {self.schema}.{self.tablename}"
+        )
+        # returning this
+        # passing through
+        return self._result

flowtask/components/CopyToBigQuery.py

@@ -0,0 +1,243 @@
+import asyncio
+from collections.abc import Callable
+import math
+import pandas as pd
+from asyncdb import AsyncDB
+from asyncdb.exceptions import (
+    StatementError,
+    DataError
+)
+from .CopyTo import CopyTo
+from ..interfaces.dataframes import PandasDataframe
+from ..exceptions import (
+    ComponentError,
+    DataNotFound
+)
+from querysource.conf import (
+    BIGQUERY_CREDENTIALS,
+    BIGQUERY_PROJECT_ID
+)
+
+
+class CopyToBigQuery(CopyTo, PandasDataframe):
+    """
+    CopyToBigQuery.
+
+    Overview
+
+        This component allows copying data into a BigQuery table,
+        using write functionality from AsyncDB BigQuery driver.
+
+    .. table:: Properties
+        :widths: auto
+
+        +--------------+----------+-----------+--------------------------------------------+
+        | Name         | Required | Summary                                                |
+        +--------------+----------+-----------+--------------------------------------------+
+        | tablename    |   Yes    | Name of the table in                                   |
+        |              |          | BigQuery                                               |
+        +--------------+----------+-----------+--------------------------------------------+
+        | schema       |   Yes    | Name of the dataset                                    |
+        |              |          | where the table is located                             |
+        +--------------+----------+-----------+--------------------------------------------+
+        | truncate     |   Yes    | This option indicates if the component should empty    |
+        |              |          | before copying the new data to the table. If set to    |
+        |              |          | true, the table will be truncated before saving data.  |
+        +--------------+----------+-----------+--------------------------------------------+
+        | use_buffer   |   No     | When activated, this option allows optimizing the      |
+        |              |          | performance of the task when dealing with large        |
+        |              |          | volumes of data.                                       |
+        +--------------+----------+-----------+--------------------------------------------+
+        | credentials  |   No     | Path to BigQuery credentials JSON file                 |
+        |              |          |                                                        |
+        +--------------+----------+-----------+--------------------------------------------+
+        | project_id   |   No     | Google Cloud Project ID                                |
+        |              |          |                                                        |
+        +--------------+----------+-----------+--------------------------------------------+
+    
+
+        Example:
+
+        ```yaml
+        CopyToBigQuery:
+          schema: hisense
+          tablename: product_availability_all
+        ```
+
+    """
+    def __init__(
+        self,
+        loop: asyncio.AbstractEventLoop = None,
+        job: Callable = None,
+        stat: Callable = None,
+        **kwargs,
+    ):
+        self.pk = []
+        self.truncate: bool = False
+        self.data = None
+        self._engine = None
+        self.tablename: str = ""
+        self.schema: str = ""  # dataset in BigQuery terminology
+        self.use_chunks = False
+        self._chunksize: int = kwargs.pop('chunksize', 10000)
+        self._connection: Callable = None
+        self._project_id: str = kwargs.pop('project_id', None)
+        try:
+            self.multi = bool(kwargs["multi"])
+            del kwargs["multi"]
+        except KeyError:
+            self.multi = False
+        super().__init__(
+            loop=loop,
+            job=job,
+            stat=stat,
+            **kwargs
+        )
+        self._driver: str = 'bigquery'
+
+    def default_connection(self):
+        """default_connection.
+
+        Default Connection to BigQuery.
+        """
+        try:
+            kwargs: dict = {
+                "credentials": BIGQUERY_CREDENTIALS,
+                "project_id": self._project_id or BIGQUERY_PROJECT_ID
+            }
+            self._connection = AsyncDB(
+                'bigquery',
+                params=kwargs,
+                loop=self._loop,
+                **kwargs
+            )
+            return self._connection
+        except Exception as err:
+            raise ComponentError(
+                f"Error configuring BigQuery Connection: {err!s}"
+            ) from err
+
+    # Function to clean invalid float values
+    def clean_floats(self, data):
+        def sanitize_value(value):
+            if isinstance(value, float) and (math.isnan(value) or math.isinf(value)):
+                return None
+            return value
+
+        if isinstance(data, dict):
+            return {k: sanitize_value(v) for k, v in data.items()}
+        elif isinstance(data, list):
+            return [self.clean_floats(item) for item in data]
+        return data
+
+    async def _create_table(self):
+        """Create a Table in BigQuery if it doesn't exist."""
+        try:
+            async with await self._connection.connection() as conn:
+                # First ensure dataset exists
+                await conn.create_dataset(self.schema)
+
+                # Infer schema from DataFrame
+                bq_schema = []
+                type_mapping = {
+                    'object': 'STRING',
+                    'string': 'STRING',
+                    'int64': 'INTEGER',
+                    'float64': 'FLOAT',
+                    'bool': 'BOOLEAN',
+                    'datetime64[ns]': 'TIMESTAMP',
+                    'date': 'DATE'
+                }
+
+                for column, dtype in self.data.dtypes.items():
+                    bq_type = type_mapping.get(str(dtype), 'STRING')
+                    bq_schema.append({
+                        "name": column,
+                        "type": bq_type,
+                        "mode": "NULLABLE"
+                    })
+
+                # If create_table has clustering fields, add them to the config
+                table_config = {}
+                if hasattr(self, 'create_table'):
+                    if isinstance(self.create_table, dict) and 'pk' in self.create_table:
+                        table_config['clustering_fields'] = self.create_table['pk']
+
+                # Create the table with the inferred schema
+                await conn.create_table(
+                    table_id=self.tablename,
+                    dataset_id=self.schema,
+                    schema=bq_schema,
+                    **table_config
+                )
+
+        except Exception as err:
+            raise ComponentError(
+                f"Error creating BigQuery table: {err}"
+            ) from err
+
+    async def _truncate_table(self):
+        """Truncate the BigQuery table using the driver's built-in method."""
+        async with await self._connection.connection() as conn:
+            await self._connection.truncate_table(
+                table_id=self.tablename,
+                dataset_id=self.schema
+            )
+
+    async def _copy_dataframe(self):
+        """Copy a pandas DataFrame to BigQuery."""
+        try:
+            # Clean NA/NaT values from string fields
+            str_cols = self.data.select_dtypes(include=["string"])
+            if not str_cols.empty:
+                self.data[str_cols.columns] = str_cols.astype(object).where(
+                    pd.notnull(str_cols), None
+                )
+
+            # Clean datetime fields - convert to ISO format strings
+            datetime_types = ["datetime64", "datetime64[ns]"]
+            datetime_cols = self.data.select_dtypes(include=datetime_types)
+            if not datetime_cols.empty:
+                for col in datetime_cols.columns:
+                    # self.data[col] = self.data[col].apply(
+                    #     lambda x: x.isoformat() if pd.notnull(x) else None
+                    # )
+                    self.data[col] = self.data[col].dt.tz_localize(None)
+            async with await self._connection.connection() as conn:
+                result = await conn.write(
+                    data=self.data,
+                    table_id=self.tablename,
+                    dataset_id=self.schema,
+                    use_pandas=True,
+                    if_exists="append"
+                )
+                print('CopyTo: ', result)
+        except StatementError as err:
+            raise ComponentError(
+                f"Statement error: {err}"
+            ) from err
+        except DataError as err:
+            raise ComponentError(
+                f"Data error: {err}"
+            ) from err
+        except Exception as err:
+            raise ComponentError(
+                f"{self.StepName} Error: {err!s}"
+            ) from err
+
+    async def _copy_iterable(self):
+        """Copy an iterable to BigQuery."""
+        try:
+            async with await self._connection.connection() as conn:
+                await conn.write(
+                    data=self.data,
+                    table_id=self.tablename,
+                    dataset_id=self.schema,
+                    use_pandas=False,
+                    if_exists="append",
+                    batch_size=self._chunksize
+                )
+        except Exception as err:
+            raise ComponentError(
+                f"Error copying iterable to BigQuery: {err}"
+            ) from err