sdk-seshat-python 0.3.14__tar.gz → 0.3.16__tar.gz

{sdk_seshat_python-0.3.14 → sdk_seshat_python-0.3.16}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sdk-seshat-python
-Version: 0.3.14
+Version: 0.3.16
 Summary: Seshat python SDK is a library to help create ML data pipelines.
 License: Commercial - see LICENSE.txt
 Author: SeshatLabs
@@ -25,6 +25,7 @@ Requires-Dist: loguru (>=0.7.3,<0.8.0)
 Requires-Dist: memory-profiler (>=0.61.0,<0.62.0)
 Requires-Dist: openai (>=1.73.0,<2.0.0)
 Requires-Dist: pandas (>=2.2.1,<3.0.0)
+Requires-Dist: psycopg2 (>=2.9,<3.0) ; extra == "postgres-support"
 Requires-Dist: pyarmor (>=8.5.1,<9.0.0)
 Requires-Dist: pydantic (>=2.7.4,<3.0.0)
 Requires-Dist: pyspark (>=3.5.1,<4.0.0)

{sdk_seshat_python-0.3.14 → sdk_seshat_python-0.3.16}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "sdk-seshat-python"
-version = "0.3.14"
+version = "0.3.16"
 description = "Seshat python SDK is a library to help create ML data pipelines."
 authors =  ["SeshatLabs <info@seshatlabs.xyz>"]
 packages = [{ include = "seshat", from = "." }]
@@ -33,11 +33,12 @@ langchain-community = "^0.3.21"
 langchain-openai = "^0.3.12"
 pyarmor = "^8.5.1"
 croniter = "^6.0.0"
+psycopg2 = { version = "^2.9", optional = true }
 python-logstash-async = "^4.0.2"
 
 [tool.poetry.extras]
 flipside_support = ["flipside"]
-postgres_support = ["psycopg2-binary"]
+postgres_support = ["psycopg2"]
 
 [tool.poetry.group.dev.dependencies]
 flake8 = "^7.0.0"

{sdk_seshat_python-0.3.14 → sdk_seshat_python-0.3.16}/seshat/source/exceptions.py

@@ -14,3 +14,7 @@ class FlipSideApiError(Exception):
     def __init__(self, e):
         message = f"An api error occurred: {str(e)}"
         super().__init__(message)
+
+
+class PostgresConstraintError(Exception):
+    pass

{sdk_seshat_python-0.3.14 → sdk_seshat_python-0.3.16}/seshat/source/mixins.py

@@ -1,7 +1,15 @@
-from typing import Optional, Callable
+from typing import Callable, Optional
 
 import sqlalchemy as db
-from sqlalchemy import Engine, create_engine, inspect, Column, MetaData, Table
+from sqlalchemy import (
+    Column,
+    Engine,
+    MetaData,
+    PrimaryKeyConstraint,
+    Table,
+    create_engine,
+    inspect,
+)
 
 from seshat.transformer.schema import Schema
 
@@ -78,11 +86,21 @@ class SQLMixin:
 
     def create_table(self, schema: Schema, table: str):
         table_columns = []
+        pk_cols = []
         for col in schema.cols:
             col_name = col.to or col.original
             col_type = getattr(db, col.dtype or "String")
-            table_columns.append(Column(col_name, col_type))
-        _, metadata = self.get_table(table, False, *table_columns, extend_existing=True)
+            table_columns.append(Column(col_name, col_type, primary_key=col.is_id))
+            if col.is_id:
+                pk_cols.append(col_name)
+        constraints = []
+        if pk_cols:
+            constraints.append(
+                PrimaryKeyConstraint(*pk_cols, name=f"{table}_pk_{'_'.join(pk_cols)}")
+            )
+        _, metadata = self.get_table(
+            table, False, *table_columns, *constraints, extend_existing=True
+        )
         metadata.create_all(self.get_engine())
 
     def get_table(self, table_name, autoload, *args, **kwargs):

{sdk_seshat_python-0.3.14 → sdk_seshat_python-0.3.16}/seshat/source/saver/database.py

@@ -2,27 +2,43 @@ import hashlib
 import statistics
 from typing import List
 
-from sqlalchemy import Index
-from sqlalchemy import inspect, select, update, and_, true
+import sqlalchemy as db
+from sqlalchemy import (
+    Column,
+    Index,
+    MetaData,
+    Table,
+    and_,
+    inspect,
+    select,
+    true,
+    update,
+)
 
 from seshat.data_class import SFrame
 from seshat.general.exceptions import InvalidArgumentsError
 from seshat.source.mixins import SQLMixin
 from seshat.source.saver import Saver
-from seshat.source.saver.base import SaveConfig
+from seshat.source.saver.base import POSTGRES, SaveConfig
 from seshat.source.saver.utils import PostgresUtils
+from seshat.transformer.schema import Schema
 from seshat.transformer.schema.base import UpdateFuncs
+from seshat.transformer.trimmer.base import NaNTrimmer
 
 
 class SQLDBSaver(SQLMixin, Saver):
     def save(self, sf: SFrame, *args, **kwargs):
+
         for config in self.save_configs:
             self.fill_cols_to_field(config.schema)
             self.ensure_table_exists(config.table, config.schema)
             self.create_index(config)
+
+            selected_sf = sf.get(config.sf_key)
+            selected_sf = self.drop_nan_ids(selected_sf, config.schema)
+
             if config.clear_table:
                 self.delete(config.table)
-            selected_sf = sf.get(config.sf_key)
             if config.strategy == "update":
                 self.update(selected_sf, config)
             elif config.strategy == "copy":
@@ -30,15 +46,25 @@ class SQLDBSaver(SQLMixin, Saver):
             else:
                 self.insert(selected_sf, config)
 
+    def ensure_table_exists(self, table: str, schema: Schema):
+        engine = self.get_engine()
+        if table in inspect(engine).get_table_names():
+            return
+        self.create_table(schema, table)
+
+    def create_table(self, schema: Schema, table: str):
+        table_columns = []
+        for col in schema.cols:
+            col_name = col.to
+            col_type = getattr(db, col.dtype or "String")
+            table_columns.append(Column(col_name, col_type))
+        _, metadata = self.get_table(table, False, *table_columns, extend_existing=True)
+        metadata.create_all(self.get_engine())
+
     def delete(self, table_name):
         table, _ = self.get_table(table_name, autoload=True)
         self.write_on_db(table.delete())
 
-    def drop_table(self, table_name):
-        if table_name in inspect(self.get_engine()).get_table_names():
-            table, _ = self.get_table(table_name, autoload=True)
-            table.drop(self.get_engine())
-
     def insert(self, selected_sf: SFrame, config: SaveConfig):
         values = self.prepare_sf_to_insert(selected_sf, config).to_dict()
         table, _ = self.get_table(config.table, autoload=True)
@@ -51,7 +77,18 @@ class SQLDBSaver(SQLMixin, Saver):
     def update(self, selected_sf: SFrame, config: SaveConfig):
         table, _ = self.get_table(config.table, autoload=True)
         values = selected_sf.to_dict()
-        id_cols = config.schema.get_id(return_first=False)
+        id_cols_schema = config.schema.get_id(return_first=False)
+
+        if self.db_type == POSTGRES:
+            try:
+                PostgresUtils.ensure_unique_constraint(
+                    self.get_engine(), table, config.schema
+                )
+                stmt = PostgresUtils.generate_upsert_stmt(table, values, config)
+                self.write_on_db(stmt)
+                return
+            except Exception:
+                pass
 
         rows_to_update = self._get_existing_rows(table, values, config.schema)
         rows_to_create = []
@@ -59,7 +96,7 @@ class SQLDBSaver(SQLMixin, Saver):
 
         if rows_to_create:
             self.write_on_db(table.insert(), rows_to_create)
-        db_id_cols = tuple(getattr(table.c, id_col.to) for id_col in id_cols)
+        db_id_cols = tuple(getattr(table.c, id_col.to) for id_col in id_cols_schema)
         for row in rows_to_update:
             condition = None
             for db_id in db_id_cols:
@@ -68,7 +105,6 @@ class SQLDBSaver(SQLMixin, Saver):
                     condition = new_condition
                 else:
                     condition = and_(condition, new_condition)
-
             update_query = update(table).where(condition).values(row)
             self.write_on_db(update_query)
 
@@ -137,6 +173,12 @@ class SQLDBSaver(SQLMixin, Saver):
         )
         return self.get_from_db(query)
 
+    def get_table(self, table_name, autoload, *args, **kwargs):
+        metadata = MetaData()
+        if autoload:
+            kwargs.setdefault("autoload_with", self.get_engine())
+        return Table(table_name, metadata, *args, **kwargs), metadata
+
     def get_from_db(self, query):
         with self.get_engine().connect() as conn:
             result = conn.execute(query)
@@ -177,5 +219,15 @@ class SQLDBSaver(SQLMixin, Saver):
         joined_cols = "-".join(set(columns))
         return hashlib.sha256(joined_cols.encode("utf-8")).hexdigest()
 
+    def drop_nan_ids(self, sf: SFrame, schema: Schema):
+        seen = set()
+        for col in schema.get_id(return_first=False):
+            if col.original in seen:
+                continue
+            seen.add(col.original)
+            trimmer = NaNTrimmer(subset=[col.original])
+            sf = trimmer(sf)
+        return sf
+
     def calculate_complexity(self):
         return 60

sdk_seshat_python-0.3.16/seshat/source/saver/utils/postgres.py

@@ -0,0 +1,105 @@
+import io
+
+import pandas as pd
+from sqlalchemy import func, inspect, text
+from sqlalchemy.dialects import postgresql as pg
+from sqlalchemy.engine import Engine
+
+from seshat.source.exceptions import PostgresConstraintError
+
+COPY_QUERY = (
+    "COPY {table_name} FROM STDIN WITH (FORMAT csv, HEADER TRUE,  DELIMITER '\t');"
+)
+
+
+class PostgresUtils:
+    @staticmethod
+    def copy(engine: Engine, table_name: str, df: pd.DataFrame):
+        csv_content = io.StringIO()
+        df.to_csv(csv_content, sep="\t", header=True, index=False)
+        csv_content.seek(0)
+        conn = engine.raw_connection()
+        cur = conn.cursor()
+        cur.copy_expert(COPY_QUERY.format(table_name=table_name), csv_content)
+        conn.commit()
+        cur.close()
+        conn.close()
+
+    @staticmethod
+    def get_constraint_name(table_name, id_columns):
+        # Generate a unique constraint name based on table and id columns
+        cols_part = "_".join(sorted(id_columns))
+        return f"uq_{table_name}_{cols_part}"
+
+    @staticmethod
+    def constraint_exists(engine, table, constraint_name):
+        # Check if a unique constraint exists on a table
+        insp = inspect(engine)
+        for cons in insp.get_unique_constraints(table.name):
+            if cons["name"] == constraint_name:
+                return True
+        return False
+
+    @staticmethod
+    def ensure_unique_constraint(engine, table, schema):
+        # Ensure a unique constraint exists for the given table and schema
+        id_cols = [col.to for col in schema.get_id(return_first=False)]
+        constraint_name = PostgresUtils.get_constraint_name(table.name, id_cols)
+        if PostgresUtils.constraint_exists(engine, table, constraint_name):
+            return True
+        try:
+            sql = 'ALTER TABLE "{table_name}" ADD CONSTRAINT "{constraint_name}" UNIQUE ({cols});'.format(
+                table_name=table.name,
+                constraint_name=constraint_name,
+                cols=", ".join(['"{}"'.format(col) for col in id_cols]),
+            )
+            with engine.connect() as conn:
+                conn.execute(text(sql))
+            return True
+        except Exception as e:
+            if "duplicate" in str(e).lower() or "unique" in str(e).lower():
+                raise PostgresConstraintError(
+                    f"Cannot create unique constraint '{constraint_name}' on table '{table.name}': "
+                    f"duplicate data exists for columns {id_cols}."
+                )
+            elif "permission" in str(e).lower() or "not allowed" in str(e).lower():
+                raise PostgresConstraintError(
+                    f"Insufficient privileges to create unique constraint '{constraint_name}' on table '{table.name}'."
+                )
+            else:
+                raise
+
+    @staticmethod
+    def generate_upsert_stmt(table, values, config):
+        # Perform an upsert (insert or update) on a Postgres table with custom update functions
+        id_cols = config.schema.get_id(return_first=False)
+        update_cols = [col for col in config.schema.cols if not col.is_id]
+        col_to_func_map = {
+            col.original: col.update_func or "replace"
+            for col in config.schema.cols
+            if not col.is_id
+        }
+        upsert_rows = []
+        for row in values:
+            upsert_row = {col.to: row[col.original] for col in config.schema.cols}
+            upsert_rows.append(upsert_row)
+        insert_stmt = pg.insert(table)
+        update_dict = {}
+        for col in update_cols:
+            func_name = col_to_func_map.get(col.original)
+            if func_name == "sum":
+                update_dict[col.to] = func.coalesce(table.c[col.to], 0) + func.coalesce(
+                    insert_stmt.excluded[col.to], 0
+                )
+            elif func_name == "mean":
+                update_dict[col.to] = (
+                    func.coalesce(table.c[col.to], 0)
+                    + func.coalesce(insert_stmt.excluded[col.to], 0)
+                ) / 2
+            else:
+                update_dict[col.to] = insert_stmt.excluded[col.to]
+        conflict_cols = [col.to for col in id_cols]
+        stmt = insert_stmt.values(upsert_rows)
+        return stmt.on_conflict_do_update(
+            index_elements=conflict_cols, set_=update_dict
+        )

sdk_seshat_python-0.3.14/seshat/source/saver/utils/postgres.py

@@ -1,22 +0,0 @@
-import io
-
-import pandas as pd
-from sqlalchemy.engine import Engine
-
-COPY_QUERY = (
-    "COPY {table_name} FROM STDIN WITH (FORMAT csv, HEADER TRUE,  DELIMITER '\t');"
-)
-
-
-class PostgresUtils:
-    @staticmethod
-    def copy(engine: Engine, table_name: str, df: pd.DataFrame):
-        csv_content = io.StringIO()
-        df.to_csv(csv_content, sep="\t", header=True, index=False)
-        csv_content.seek(0)
-        conn = engine.raw_connection()
-        cur = conn.cursor()
-        cur.copy_expert(COPY_QUERY.format(table_name=table_name), csv_content)
-        conn.commit()
-        cur.close()
-        conn.close()