omnata-plugin-runtime 0.2.67__py3-none-any.whl → 0.2.69__py3-none-any.whl

omnata_plugin_runtime/monkey_patching.py

@@ -0,0 +1,295 @@
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+#
+
+from __future__ import annotations
+
+import collections.abc
+import os
+import warnings
+from functools import partial
+from logging import getLogger
+from tempfile import TemporaryDirectory
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Iterable,
+    Iterator,
+    Literal,
+    Sequence,
+    TypeVar,
+)
+
+from snowflake.connector import ProgrammingError
+from snowflake.connector.options import pandas
+from snowflake.connector.telemetry import TelemetryData, TelemetryField
+from snowflake.connector.util_text import random_string
+
+from snowflake.connector.connection import SnowflakeConnection
+from snowflake.connector.pandas_tools import (
+    _create_temp_file_format,
+    _create_temp_stage,
+    build_location_helper,
+    chunk_helper
+)
+
+logger = getLogger(__name__)
+
+
+def write_pandas(
+    conn: SnowflakeConnection,
+    df: pandas.DataFrame,
+    table_name: str,
+    database: str | None = None,
+    schema: str | None = None,
+    chunk_size: int | None = None,
+    compression: str = "gzip",
+    on_error: str = "abort_statement",
+    parallel: int = 4,
+    quote_identifiers: bool = True,
+    auto_create_table: bool = False,
+    create_temp_table: bool = False,
+    overwrite: bool = False,
+    table_type: Literal["", "temp", "temporary", "transient"] = "",
+    parquet_engine: Literal["auto", "pyarrow", "fastparquet"] = "auto",
+    **kwargs: Any,
+) -> tuple[
+    bool,
+    int,
+    int,
+    Sequence[
+        tuple[
+            str,
+            str,
+            int,
+            int,
+            int,
+            int,
+            str | None,
+            int | None,
+            int | None,
+            str | None,
+        ]
+    ],
+]:
+    """Allows users to most efficiently write back a pandas DataFrame to Snowflake.
+
+    It works by dumping the DataFrame into Parquet files, uploading them and finally copying their data into the table.
+
+    Returns whether all files were ingested correctly, number of chunks uploaded, and number of rows ingested
+    with all of the COPY INTO command's output for debugging purposes.
+
+        Example usage:
+            import pandas
+            from snowflake.connector.pandas_tools import write_pandas
+
+            df = pandas.DataFrame([('Mark', 10), ('Luke', 20)], columns=['name', 'balance'])
+            success, nchunks, nrows, _ = write_pandas(cnx, df, 'customers')
+
+    Args:
+        conn: Connection to be used to communicate with Snowflake.
+        df: Dataframe we'd like to write back.
+        table_name: Table name where we want to insert into.
+        database: Database schema and table is in, if not provided the default one will be used (Default value = None).
+        schema: Schema table is in, if not provided the default one will be used (Default value = None).
+        chunk_size: Number of elements to be inserted once, if not provided all elements will be dumped once
+            (Default value = None).
+        compression: The compression used on the Parquet files, can only be gzip, or snappy. Gzip gives supposedly a
+            better compression, while snappy is faster. Use whichever is more appropriate (Default value = 'gzip').
+        on_error: Action to take when COPY INTO statements fail, default follows documentation at:
+            https://docs.snowflake.com/en/sql-reference/sql/copy-into-table.html#copy-options-copyoptions
+            (Default value = 'abort_statement').
+        parallel: Number of threads to be used when uploading chunks, default follows documentation at:
+            https://docs.snowflake.com/en/sql-reference/sql/put.html#optional-parameters (Default value = 4).
+        quote_identifiers: By default, identifiers, specifically database, schema, table and column names
+            (from df.columns) will be quoted. If set to False, identifiers are passed on to Snowflake without quoting.
+            I.e. identifiers will be coerced to uppercase by Snowflake.  (Default value = True)
+        auto_create_table: When true, will automatically create a table with corresponding columns for each column in
+            the passed in DataFrame. The table will not be created if it already exists
+        create_temp_table: (Deprecated) Will make the auto-created table as a temporary table
+        overwrite: When true, and if auto_create_table is true, then it drops the table. Otherwise, it
+        truncates the table. In both cases it will replace the existing contents of the table with that of the passed in
+            Pandas DataFrame.
+        table_type: The table type of to-be-created table. The supported table types include ``temp``/``temporary``
+            and ``transient``. Empty means permanent table as per SQL convention.
+
+    Returns:
+        Returns the COPY INTO command's results to verify ingestion in the form of a tuple of whether all chunks were
+        ingested correctly, # of chunks, # of ingested rows, and ingest's output.
+    """
+    if database is not None and schema is None:
+        raise ProgrammingError(
+            "Schema has to be provided to write_pandas when a database is provided"
+        )
+    # This dictionary maps the compression algorithm to Snowflake put copy into command type
+    # https://docs.snowflake.com/en/sql-reference/sql/copy-into-table.html#type-parquet
+    compression_map = {"gzip": "auto", "snappy": "snappy"}
+    if compression not in compression_map.keys():
+        raise ProgrammingError(
+            f"Invalid compression '{compression}', only acceptable values are: {compression_map.keys()}"
+        )
+
+    if create_temp_table:
+        warnings.warn(
+            "create_temp_table is deprecated, we still respect this parameter when it is True but "
+            'please consider using `table_type="temp"` instead',
+            DeprecationWarning,
+            # warnings.warn -> write_pandas
+            stacklevel=2,
+        )
+        table_type = "temp"
+
+    if table_type and table_type.lower() not in ["temp", "temporary", "transient"]:
+        raise ValueError(
+            "Unsupported table type. Expected table types: temp/temporary, transient"
+        )
+
+    if chunk_size is None:
+        chunk_size = len(df)
+
+    if not (
+        isinstance(df.index, pandas.RangeIndex)
+        and 1 == df.index.step
+        and 0 == df.index.start
+    ):
+        warnings.warn(
+            f"Pandas Dataframe has non-standard index of type {str(type(df.index))} which will not be written."
+            f" Consider changing the index to pd.RangeIndex(start=0,...,step=1) or "
+            f"call reset_index() to keep index as column(s)",
+            UserWarning,
+            stacklevel=2,
+        )
+
+    cursor = conn.cursor()
+    stage_location = _create_temp_stage(
+        cursor,
+        database,
+        schema,
+        quote_identifiers,
+        compression,
+        auto_create_table,
+        overwrite,
+    )
+
+    with TemporaryDirectory() as tmp_folder:
+        for i, chunk in chunk_helper(df, chunk_size):
+            chunk_path = os.path.join(tmp_folder, f"file{i}.txt")
+            # Dump chunk into parquet file
+            chunk.to_parquet(chunk_path, compression=compression, engine=parquet_engine, **kwargs)
+            # Upload parquet file
+            upload_sql = (
+                "PUT /* Python:snowflake.connector.pandas_tools.write_pandas() */ "
+                "'file://{path}' @{stage_location} PARALLEL={parallel}"
+            ).format(
+                path=chunk_path.replace("\\", "\\\\").replace("'", "\\'"),
+                stage_location=stage_location,
+                parallel=parallel,
+            )
+            logger.debug(f"uploading files with '{upload_sql}'")
+            cursor.execute(upload_sql, _is_internal=True)
+            # Remove chunk file
+            os.remove(chunk_path)
+
+    # in Snowflake, all parquet data is stored in a single column, $1, so we must select columns explicitly
+    # see (https://docs.snowflake.com/en/user-guide/script-data-load-transform-parquet.html)
+    if quote_identifiers:
+        quote = '"'
+        # if the column name contains a double quote, we need to escape it by replacing with two double quotes
+        # https://docs.snowflake.com/en/sql-reference/identifiers-syntax#double-quoted-identifiers
+        snowflake_column_names = [str(c).replace('"', '""') for c in df.columns]
+    else:
+        quote = ""
+        snowflake_column_names = list(df.columns)
+    columns = quote + f"{quote},{quote}".join(snowflake_column_names) + quote
+
+    def drop_object(name: str, object_type: str) -> None:
+        drop_sql = f"DROP {object_type.upper()} IF EXISTS {name} /* Python:snowflake.connector.pandas_tools.write_pandas() */"
+        logger.debug(f"dropping {object_type} with '{drop_sql}'")
+        cursor.execute(drop_sql, _is_internal=True)
+
+    if auto_create_table or overwrite:
+        file_format_location = _create_temp_file_format(
+            cursor, database, schema, quote_identifiers, compression_map[compression]
+        )
+        infer_schema_sql = f"SELECT COLUMN_NAME, TYPE FROM table(infer_schema(location=>'@{stage_location}', file_format=>'{file_format_location}'))"
+        logger.debug(f"inferring schema with '{infer_schema_sql}'")
+        column_type_mapping = dict(
+            cursor.execute(infer_schema_sql, _is_internal=True).fetchall()
+        )
+        # Infer schema can return the columns out of order depending on the chunking we do when uploading
+        # so we have to iterate through the dataframe columns to make sure we create the table with its
+        # columns in order
+        create_table_columns = ", ".join(
+            [
+                f"{quote}{snowflake_col}{quote} {column_type_mapping[col]}"
+                for snowflake_col, col in zip(snowflake_column_names, df.columns)
+            ]
+        )
+
+        target_table_location = build_location_helper(
+            database,
+            schema,
+            random_string() if overwrite else table_name,
+            quote_identifiers,
+        )
+
+        create_table_sql = (
+            f"CREATE {table_type.upper()} TABLE IF NOT EXISTS {target_table_location} "
+            f"({create_table_columns})"
+            f" /* Python:snowflake.connector.pandas_tools.write_pandas() */ "
+        )
+        logger.debug(f"auto creating table with '{create_table_sql}'")
+        cursor.execute(create_table_sql, _is_internal=True)
+        # need explicit casting when the underlying table schema is inferred
+        parquet_columns = "$1:" + ",$1:".join(
+            f"{quote}{snowflake_col}{quote}::{column_type_mapping[col]}"
+            for snowflake_col, col in zip(snowflake_column_names, df.columns)
+        )
+    else:
+        target_table_location = build_location_helper(
+            database=database,
+            schema=schema,
+            name=table_name,
+            quote_identifiers=quote_identifiers,
+        )
+        parquet_columns = "$1:" + ",$1:".join(
+            f"{quote}{snowflake_col}{quote}" for snowflake_col in snowflake_column_names
+        )
+
+    try:
+        copy_into_sql = (
+            f"COPY INTO {target_table_location} /* Python:snowflake.connector.pandas_tools.write_pandas() */ "
+            f"({columns}) "
+            f"FROM (SELECT {parquet_columns} FROM @{stage_location}) "
+            f"FILE_FORMAT=(TYPE=PARQUET COMPRESSION={compression_map[compression]}{' BINARY_AS_TEXT=FALSE' if auto_create_table or overwrite else ''}) "
+            f"PURGE=TRUE ON_ERROR={on_error}"
+        )
+        logger.debug(f"copying into with '{copy_into_sql}'")
+        copy_results = cursor.execute(copy_into_sql, _is_internal=True).fetchall()
+
+        if overwrite:
+            original_table_location = build_location_helper(
+                database=database,
+                schema=schema,
+                name=table_name,
+                quote_identifiers=quote_identifiers,
+            )
+            drop_object(original_table_location, "table")
+            rename_table_sql = f"ALTER TABLE {target_table_location} RENAME TO {original_table_location} /* Python:snowflake.connector.pandas_tools.write_pandas() */"
+            logger.debug(f"rename table with '{rename_table_sql}'")
+            cursor.execute(rename_table_sql, _is_internal=True)
+    except ProgrammingError:
+        if overwrite:
+            drop_object(target_table_location, "table")
+        raise
+    finally:
+        cursor._log_telemetry_job_data(TelemetryField.PANDAS_WRITE, TelemetryData.TRUE)
+        cursor.close()
+
+    return (
+        all(e[1] == "LOADED" for e in copy_results),
+        len(copy_results),
+        sum(int(e[3]) for e in copy_results),
+        copy_results,
+    )

omnata_plugin_runtime/omnata_plugin.py

@@ -67,6 +67,9 @@ from .rate_limiting import (
 logger = getLogger(__name__)
 SortDirectionType = Literal["asc", "desc"]
 
+import snowflake.connector.pandas_tools
+from .monkey_patching import write_pandas
+snowflake.connector.pandas_tools.write_pandas = write_pandas
 
 class PluginManifest(SubscriptableBaseModel):
     """
@@ -650,6 +653,7 @@ class OutboundSyncRequest(SyncRequest):
                 quote_identifiers=False,
                 table_name=self._full_results_table_name,
                 auto_create_table=False,
+                engine='fastparquet'
             )
             if not success:
                 raise ValueError(
@@ -1081,6 +1085,8 @@ class InboundSyncRequest(SyncRequest):
                     "RETRIEVE_DATE",
                 ],
             )
+        # provide record data as a string so that parquet writing doesn't get tangled up trying to guess types
+        results_df["RECORD_DATA"] = results_df["RECORD_DATA"].astype(str)
         # trim out the columns we don't need to return
         return results_df[
             results_df.columns.intersection(

{omnata_plugin_runtime-0.2.67.dist-info → omnata_plugin_runtime-0.2.69.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: omnata-plugin-runtime
-Version: 0.2.67
+Version: 0.2.69
 Summary: Classes and common runtime components for building and running Omnata Plugins
 Author: James Weakley
 Author-email: james.weakley@omnata.com

{omnata_plugin_runtime-0.2.67.dist-info → omnata_plugin_runtime-0.2.69.dist-info}/RECORD

@@ -3,10 +3,11 @@ omnata_plugin_runtime/api.py,sha256=vKq7goVPX5cPQ9CVN9l8RmpJDcqDwS5y9v1IWhWjBbk,
 omnata_plugin_runtime/configuration.py,sha256=NICse7qMtReIYY4ZCu8ng4QDJ4rFP0g3mZwc8m1Xl54,32247
 omnata_plugin_runtime/forms.py,sha256=_KqSMQG749wImLKxPZh3B3doTZMbP5jDvF6BhQNkPCM,17375
 omnata_plugin_runtime/logging.py,sha256=Q6eSqrr3SzwfVAg4r4sV1dlxeNS_PzOtZfieoWUEOZQ,3232
-omnata_plugin_runtime/omnata_plugin.py,sha256=0oqCFj4IGt_q-_nJSjk5QtQuK1BNcec2J9PmyzneD24,85265
+omnata_plugin_runtime/monkey_patching.py,sha256=ddYEdVBR_SlijNuKLGSRxxEboiOMcDjQ9lKv4KrahVc,12637
+omnata_plugin_runtime/omnata_plugin.py,sha256=VGTmRr4XYitkLMBPzHswKHUg_Ie0eFiEOUjwH1cXMRQ,85630
 omnata_plugin_runtime/plugin_entrypoints.py,sha256=dV_JOEWffdkAkofFYJMAsM_-gE9OpM2-paI2dsh5Rzo,24352
 omnata_plugin_runtime/rate_limiting.py,sha256=QO8VB1H8al6a8-ydohUmL0c5JynXG2bulmuPRs2-2-Y,14910
-omnata_plugin_runtime-0.2.67.dist-info/LICENSE,sha256=IMF9i4xIpgCADf0U-V1cuf9HBmqWQd3qtI3FSuyW4zE,26526
-omnata_plugin_runtime-0.2.67.dist-info/METADATA,sha256=U0-X5qtmLZgJobz2xNO4dCRpPY0YGdja_sRNBeFXrPo,1086
-omnata_plugin_runtime-0.2.67.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
-omnata_plugin_runtime-0.2.67.dist-info/RECORD,,
+omnata_plugin_runtime-0.2.69.dist-info/LICENSE,sha256=IMF9i4xIpgCADf0U-V1cuf9HBmqWQd3qtI3FSuyW4zE,26526
+omnata_plugin_runtime-0.2.69.dist-info/METADATA,sha256=-dAiaEMQ2YOx1TuL9BHNM8uiJm9JXOrNGti4FNtere4,1086
+omnata_plugin_runtime-0.2.69.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
+omnata_plugin_runtime-0.2.69.dist-info/RECORD,,