scitex 2.0.0__py2.py3-none-any.whl

scitex/db/_PostgreSQLMixins/_BackupMixin.py

@@ -0,0 +1,166 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Time-stamp: "2024-11-24 23:06:03 (ywatanabe)"
+# File: ./scitex_repo/src/scitex/db/_PostgreSQLMixins/_BackupMixin.py
+
+THIS_FILE = (
+    "/home/ywatanabe/proj/scitex_repo/src/scitex/db/_PostgreSQLMixins/_BackupMixin.py"
+)
+
+import subprocess
+import os
+from typing import Optional, List
+from .._BaseMixins._BaseBackupMixin import _BaseBackupMixin
+
+
+class _BackupMixin(_BaseBackupMixin):
+    def backup_table(self, table: str, file_path: str) -> None:
+        """Backup a specific table using pg_dump"""
+        connection_params = self._get_connection_params()
+        cmd = [
+            "pg_dump",
+            "-h",
+            connection_params["host"],
+            "-p",
+            str(connection_params["port"]),
+            "-U",
+            connection_params["user"],
+            "-d",
+            connection_params["database"],
+            "-t",
+            table,
+            "-f",
+            file_path,
+        ]
+
+        env = os.environ.copy()
+        env["PGPASSWORD"] = connection_params["password"]
+
+        try:
+            subprocess.run(cmd, env=env, check=True, capture_output=True)
+        except subprocess.CalledProcessError as e:
+            raise Exception(f"Backup failed: {e.stderr.decode()}")
+
+    def restore_table(self, table: str, file_path: str) -> None:
+        """Restore a specific table from backup"""
+        self._check_writable()
+        connection_params = self._get_connection_params()
+        cmd = [
+            "psql",
+            "-h",
+            connection_params["host"],
+            "-p",
+            str(connection_params["port"]),
+            "-U",
+            connection_params["user"],
+            "-d",
+            connection_params["database"],
+            "-f",
+            file_path,
+        ]
+
+        env = os.environ.copy()
+        env["PGPASSWORD"] = connection_params["password"]
+
+        try:
+            subprocess.run(cmd, env=env, check=True, capture_output=True)
+        except subprocess.CalledProcessError as e:
+            raise Exception(f"Restore failed: {e.stderr.decode()}")
+
+    def backup_database(self, file_path: str) -> None:
+        """Backup entire database using pg_dump"""
+        connection_params = self._get_connection_params()
+        cmd = [
+            "pg_dump",
+            "-h",
+            connection_params["host"],
+            "-p",
+            str(connection_params["port"]),
+            "-U",
+            connection_params["user"],
+            "-d",
+            connection_params["database"],
+            "-F",
+            "c",  # Custom format
+            "-f",
+            file_path,
+        ]
+
+        env = os.environ.copy()
+        env["PGPASSWORD"] = connection_params["password"]
+
+        try:
+            subprocess.run(cmd, env=env, check=True, capture_output=True)
+        except subprocess.CalledProcessError as e:
+            raise Exception(f"Database backup failed: {e.stderr.decode()}")
+
+    def restore_database(self, file_path: str) -> None:
+        """Restore entire database from backup"""
+        self._check_writable()
+        connection_params = self._get_connection_params()
+        cmd = [
+            "pg_restore",
+            "-h",
+            connection_params["host"],
+            "-p",
+            str(connection_params["port"]),
+            "-U",
+            connection_params["user"],
+            "-d",
+            connection_params["database"],
+            "--clean",  # Clean (drop) database objects before recreating
+            file_path,
+        ]
+
+        env = os.environ.copy()
+        env["PGPASSWORD"] = connection_params["password"]
+
+        try:
+            subprocess.run(cmd, env=env, check=True, capture_output=True)
+        except subprocess.CalledProcessError as e:
+            raise Exception(f"Database restore failed: {e.stderr.decode()}")
+
+    def copy_table(
+        self, source_table: str, target_table: str, where: Optional[str] = None
+    ) -> None:
+        """Copy data from one table to another"""
+        self._check_writable()
+
+        # Get column names from source table
+        self.execute(
+            f"""
+            SELECT column_name
+            FROM information_schema.columns
+            WHERE table_name = %s
+            ORDER BY ordinal_position
+        """,
+            (source_table,),
+        )
+        columns: List[str] = [row[0] for row in self.cursor.fetchall()]
+        columns_str = ", ".join(columns)
+
+        # Create the target table with the same structure
+        self.execute(
+            f"CREATE TABLE IF NOT EXISTS {target_table} AS SELECT * FROM {source_table} WHERE 1=0"
+        )
+
+        # Copy data
+        query = f"INSERT INTO {target_table} ({columns_str}) SELECT {columns_str} FROM {source_table}"
+        if where:
+            query += f" WHERE {where}"
+
+        self.execute(query)
+
+    def _get_connection_params(self) -> dict:
+        """Extract connection parameters from the current connection"""
+        dsn = self.conn.get_dsn_parameters()
+        return {
+            "host": dsn.get("host", "localhost"),
+            "port": dsn.get("port", 5432),
+            "user": dsn.get("user", ""),
+            "password": dsn.get("password", ""),
+            "database": dsn.get("dbname", ""),
+        }
+
+
+# EOF

scitex/db/_PostgreSQLMixins/_BatchMixin.py

@@ -0,0 +1,82 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Timestamp: "2025-02-27 22:14:16 (ywatanabe)"
+# File: /home/ywatanabe/proj/scitex_dev/src/scitex/db/_PostgreSQLMixins/_BatchMixin.py
+
+THIS_FILE = (
+    "/home/ywatanabe/proj/scitex_repo/src/scitex/db/_PostgreSQLMixins/_BatchMixin.py"
+)
+
+from typing import List, Any, Optional, Dict, Union
+import pandas as pd
+from .._BaseMixins._BaseBatchMixin import _BaseBatchMixin
+
+
+class _BatchMixin(_BaseBatchMixin):
+    def insert_many(
+        self,
+        table: str,
+        records: List[Dict[str, Any]],
+        batch_size: Optional[int] = None,
+    ) -> None:
+        if not records:
+            return
+
+        query = self._prepare_insert_query(table, records[0])
+        if batch_size and batch_size > 0:
+            for i in range(0, len(records), batch_size):
+                batch = records[i : i + batch_size]
+                parameters = self._prepare_batch_parameters(batch)
+                self.executemany(query, parameters)
+        else:
+            parameters = self._prepare_batch_parameters(records)
+            self.executemany(query, parameters)
+
+    def _prepare_insert_query(self, table: str, record: Dict[str, Any]) -> str:
+        columns = list(record.keys())
+        placeholders = ["%s"] * len(columns)
+        columns_str = ", ".join(columns)
+        placeholders_str = ", ".join(placeholders)
+        return f"INSERT INTO {table} ({columns_str}) VALUES ({placeholders_str})"
+
+    def _prepare_batch_parameters(self, records: List[Dict[str, Any]]) -> List[tuple]:
+        if not records:
+            return []
+
+        columns = list(records[0].keys())
+        return [tuple(record[col] for col in columns) for record in records]
+
+    def dataframe_to_sql(
+        self, df: pd.DataFrame, table: str, if_exists: str = "fail"
+    ) -> None:
+        if if_exists not in ["fail", "replace", "append"]:
+            raise ValueError("if_exists must be one of 'fail', 'replace', or 'append'")
+
+        if if_exists == "replace":
+            self.execute(f"DROP TABLE IF EXISTS {table}")
+            # Create table based on DataFrame schema
+            columns = []
+            for col, dtype in df.dtypes.items():
+                pg_type = self._map_dtype_to_postgres(dtype)
+                columns.append(f"{col} {pg_type}")
+            columns_str = ", ".join(columns)
+            self.execute(f"CREATE TABLE {table} ({columns_str})")
+
+        records = df.to_dict("records")
+        self.insert_many(table, records)
+
+    def _map_dtype_to_postgres(self, dtype) -> str:
+        dtype_str = str(dtype)
+        if "int" in dtype_str:
+            return "INTEGER"
+        elif "float" in dtype_str:
+            return "REAL"
+        elif "datetime" in dtype_str:
+            return "TIMESTAMP"
+        elif "bool" in dtype_str:
+            return "BOOLEAN"
+        else:
+            return "TEXT"
+
+
+# EOF

scitex/db/_PostgreSQLMixins/_BlobMixin.py

@@ -0,0 +1,231 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Timestamp: "2025-02-27 22:14:45 (ywatanabe)"
+# File: /home/ywatanabe/proj/scitex_dev/src/scitex/db/_PostgreSQLMixins/_BlobMixin.py
+
+THIS_FILE = "/home/ywatanabe/proj/scitex_repo/src/scitex/db/_PostgreSQLMixins/_BlobMixin.py"
+
+import psycopg2
+import numpy as np
+from typing import Any as _Any
+from typing import Dict, List, Optional, Tuple, Union
+from .._BaseMixins._BaseBlobMixin import _BaseBlobMixin
+
+
+class _BlobMixin(_BaseBlobMixin):
+    """BLOB data handling functionality for PostgreSQL"""
+
+    def save_array(
+        self,
+        table_name: str,
+        data: np.ndarray,
+        column: str = "data",
+        ids: Optional[Union[int, List[int]]] = None,
+        where: str = None,
+        additional_columns: Dict[str, _Any] = None,
+        batch_size: int = 1000,
+    ) -> None:
+        with self.transaction():
+            if not isinstance(data, (np.ndarray, list)):
+                raise ValueError("Input must be a NumPy array or list of arrays")
+
+            try:
+                if ids is not None:
+                    if isinstance(ids, int):
+                        ids = [ids]
+                        data = [data]
+                    if len(ids) != len(data):
+                        raise ValueError("Length of ids must match number of arrays")
+
+                    for id_, arr in zip(ids, data):
+                        if not isinstance(arr, np.ndarray):
+                            raise ValueError(
+                                f"Element for id {id_} must be a NumPy array"
+                            )
+
+                        binary = psycopg2.Binary(arr.tobytes())
+                        # binary = Binary(arr.tobytes())
+                        columns = [
+                            column,
+                            f"{column}_dtype",
+                            f"{column}_shape",
+                        ]
+                        values = [binary, str(arr.dtype), str(arr.shape)]
+
+                        if additional_columns:
+                            columns = list(additional_columns.keys()) + columns
+                            values = list(additional_columns.values()) + values
+
+                        update_cols = [f"{col}=%s" for col in columns]
+                        query = f"UPDATE {table_name} SET {','.join(update_cols)} WHERE id=%s"
+                        values.append(id_)
+                        self.execute(query, values)
+
+                else:
+                    if not isinstance(data, np.ndarray):
+                        raise ValueError("Single input must be a NumPy array")
+
+                    binary = psycopg2.Binary(arr.tobytes())
+                    # binary = Binary(data.tobytes())
+                    columns = [column, f"{column}_dtype", f"{column}_shape"]
+                    values = [binary, str(data.dtype), str(data.shape)]
+
+                    if additional_columns:
+                        columns = list(additional_columns.keys()) + columns
+                        values = list(additional_columns.values()) + values
+
+                    if where is not None:
+                        update_cols = [f"{col}=%s" for col in columns]
+                        query = f"UPDATE {table_name} SET {','.join(update_cols)} WHERE {where}"
+                        self.execute(query, values)
+                    else:
+                        placeholders = ",".join(["%s" for _ in columns])
+                        columns_str = ",".join(columns)
+                        query = f"INSERT INTO {table_name} ({columns_str}) VALUES ({placeholders})"
+                        self.execute(query, values)
+
+            except Exception as err:
+                raise ValueError(f"Failed to save array: {err}")
+
+    def load_array(
+        self,
+        table_name: str,
+        column: str,
+        ids: Union[int, List[int], str] = "all",
+        where: str = None,
+        order_by: str = None,
+        batch_size: int = 128,
+        dtype: np.dtype = None,
+        shape: Optional[Tuple] = None,
+    ) -> Optional[np.ndarray]:
+        try:
+            if ids == "all":
+                query = f"SELECT id FROM {table_name}"
+                if where:
+                    query += f" WHERE {where}"
+                self.execute(query)
+                ids = [row[0] for row in self.cursor.fetchall()]
+            elif isinstance(ids, int):
+                ids = [ids]
+
+            id_to_data = {}
+            unique_ids = list(set(ids))
+
+            for idx in range(0, len(unique_ids), batch_size):
+                batch_ids = unique_ids[idx : idx + batch_size]
+                placeholders = ",".join(["%s" for _ in batch_ids])
+
+                try:
+                    query = f"""
+                        SELECT id, {column},
+                               {column}_dtype,
+                               {column}_shape
+                        FROM {table_name}
+                        WHERE id IN ({placeholders})
+                    """
+                    self.execute(query, batch_ids)
+                    has_metadata = True
+                except psycopg2.Error:
+                    query = f"SELECT id, {column} FROM {table_name} WHERE id IN ({placeholders})"
+                    self.execute(query, batch_ids)
+                    has_metadata = False
+
+                if where:
+                    query += f" AND {where}"
+                if order_by:
+                    query += f" ORDER BY {order_by}"
+
+                results = self.cursor.fetchall()
+                if results:
+                    for result in results:
+                        if has_metadata:
+                            id_val, blob, dtype_str, shape_str = result
+                            data = np.frombuffer(
+                                bytes(blob), dtype=np.dtype(dtype_str)
+                            ).reshape(eval(shape_str))
+                        else:
+                            id_val, blob = result
+                            data = (
+                                np.frombuffer(bytes(blob), dtype=dtype)
+                                if dtype
+                                else np.frombuffer(bytes(blob))
+                            )
+                            if shape:
+                                data = data.reshape(shape)
+                        id_to_data[id_val] = data
+
+            all_data = [id_to_data[id_val] for id_val in ids if id_val in id_to_data]
+            return np.stack(all_data, axis=0) if all_data else None
+
+        except Exception as err:
+            raise ValueError(f"Failed to load array: {err}")
+
+    def binary_to_array(
+        self,
+        binary_data,
+        dtype_str=None,
+        shape_str=None,
+        dtype=None,
+        shape=None,
+    ):
+        if binary_data is None:
+            return None
+
+        binary_data = bytes(binary_data)
+        if dtype_str and shape_str:
+            return np.frombuffer(binary_data, dtype=np.dtype(dtype_str)).reshape(
+                eval(shape_str)
+            )
+        elif dtype and shape:
+            return np.frombuffer(binary_data, dtype=dtype).reshape(shape)
+        return binary_data
+
+    def get_array_dict(self, df, columns=None, dtype=None, shape=None):
+        result = {}
+        if columns is None:
+            columns = [
+                col
+                for col in df.columns
+                if not (col.endswith("_dtype") or col.endswith("_shape"))
+            ]
+
+        for col in columns:
+            if f"{col}_dtype" in df.columns and f"{col}_shape" in df.columns:
+                arrays = [
+                    self.binary_to_array(
+                        row[col], row[f"{col}_dtype"], row[f"{col}_shape"]
+                    )
+                    for _, row in df.iterrows()
+                ]
+            elif dtype and shape:
+                arrays = [
+                    self.binary_to_array(x, dtype=dtype, shape=shape) for x in df[col]
+                ]
+            result[col] = np.stack(arrays)
+
+        return result
+
+    def decode_array_columns(self, df, columns=None, dtype=None, shape=None):
+        if columns is None:
+            columns = [
+                col
+                for col in df.columns
+                if not (col.endswith("_dtype") or col.endswith("_shape"))
+            ]
+
+        for col in columns:
+            if f"{col}_dtype" in df.columns and f"{col}_shape" in df.columns:
+                df[col] = df.apply(
+                    lambda row: self.binary_to_array(
+                        row[col], row[f"{col}_dtype"], row[f"{col}_shape"]
+                    ),
+                    axis=1,
+                )
+            elif dtype and shape:
+                df[col] = df[col].apply(
+                    lambda x: self.binary_to_array(x, dtype=dtype, shape=shape)
+                )
+        return df
+
+
+# EOF

scitex/db/_PostgreSQLMixins/_ConnectionMixin.py

@@ -0,0 +1,92 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Timestamp: "2025-02-27 22:14:52 (ywatanabe)"
+# File: /home/ywatanabe/proj/scitex_dev/src/scitex/db/_PostgreSQLMixins/_ConnectionMixin.py
+
+THIS_FILE = (
+    "/home/ywatanabe/proj/scitex_repo/src/scitex/db/_PostgreSQLMixins/_ConnectionMixin.py"
+)
+
+from typing import Any, Tuple
+import psycopg2
+
+from .._BaseMixins._BaseConnectionMixin import _BaseConnectionMixin
+
+
+class _ConnectionMixin(_BaseConnectionMixin):
+    def __init__(
+        self,
+        dbname: str,
+        user: str,
+        password: str,
+        host: str = "localhost",
+        port: int = 5432,
+    ):
+        super().__init__()
+        self.db_config = {
+            "dbname": dbname,
+            "user": user,
+            "password": password,
+            "host": host,
+            "port": port,
+        }
+        if dbname:
+            self.connect()
+
+    def connect(self) -> None:
+        if self.conn:
+            self.close()
+
+        self.conn = psycopg2.connect(**self.db_config)
+        self.cursor = self.conn.cursor()
+
+        with self.lock:
+            self.conn.autocommit = False
+            self.cursor.execute(
+                "SET SESSION CHARACTERISTICS AS TRANSACTION ISOLATION LEVEL READ COMMITTED"
+            )
+
+    def close(self) -> None:
+        if self.cursor:
+            self.cursor.close()
+        if self.conn:
+            try:
+                self.conn.close()
+            except psycopg2.Error:
+                pass
+        self.cursor = None
+        self.conn = None
+
+    def reconnect(self) -> None:
+        if self.db_config:
+            self.connect()
+        else:
+            raise ValueError("No database configuration specified for reconnection")
+
+    def execute(self, query: str, parameters: Tuple = None) -> Any:
+        """Execute a database query."""
+        if not self.cursor:
+            raise ConnectionError("Database not connected")
+
+        try:
+            self.cursor.execute(query, parameters)
+            self.conn.commit()
+            return self.cursor
+        except psycopg2.Error as err:
+            self.conn.rollback()
+            raise psycopg2.Error(f"Query execution failed: {err}")
+
+    def executemany(self, query: str, parameters: list) -> None:
+        """Execute multiple database queries."""
+        if not self.cursor:
+            raise ConnectionError("Database not connected")
+
+        try:
+            self.cursor.executemany(query, parameters)
+            self.conn.commit()
+        except psycopg2.Error as err:
+            self.conn.rollback()
+            raise psycopg2.Error(f"Batch query execution failed: {err}")
+
+
+# EOF

scitex/db/_PostgreSQLMixins/_ImportExportMixin.py

@@ -0,0 +1,59 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Timestamp: "2025-02-27 22:14:59 (ywatanabe)"
+# File: /home/ywatanabe/proj/scitex_dev/src/scitex/db/_PostgreSQLMixins/_ImportExportMixin.py
+
+THIS_FILE = (
+    "/home/ywatanabe/proj/scitex_repo/src/scitex/db/_PostgreSQLMixins/_ImportExportMixin.py"
+)
+
+import pandas as pd
+from typing import List
+import psycopg2
+from io import StringIO
+
+
+class _ImportExportMixin:
+    def load_from_csv(
+        self,
+        table_name: str,
+        csv_path: str,
+        if_exists: str = "append",
+        batch_size: int = 10_000,
+        chunk_size: int = 100_000,
+    ) -> None:
+        with self.transaction():
+            try:
+                if if_exists == "replace":
+                    self.execute(f"TRUNCATE TABLE {table_name}")
+
+                copy_sql = f"COPY {table_name} FROM STDIN WITH CSV HEADER"
+                with open(csv_path, "r") as f:
+                    self.cursor.copy_expert(sql=copy_sql, file=f)
+
+            except (Exception, psycopg2.Error) as err:
+                raise ValueError(f"Failed to import from CSV: {err}")
+
+    def save_to_csv(
+        self,
+        table_name: str,
+        output_path: str,
+        columns: List[str] = ["*"],
+        where: str = None,
+        batch_size: int = 10_000,
+    ) -> None:
+        try:
+            columns_str = ", ".join(columns) if columns != ["*"] else "*"
+            query = f"COPY (SELECT {columns_str} FROM {table_name}"
+            if where:
+                query += f" WHERE {where}"
+            query += ") TO STDOUT WITH CSV HEADER"
+
+            with open(output_path, "w") as f:
+                self.cursor.copy_expert(sql=query, file=f)
+
+        except (Exception, psycopg2.Error) as err:
+            raise ValueError(f"Failed to export to CSV: {err}")
+
+
+# EOF

scitex/db/_PostgreSQLMixins/_IndexMixin.py

@@ -0,0 +1,64 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Timestamp: "2025-02-27 22:15:05 (ywatanabe)"
+# File: /home/ywatanabe/proj/scitex_dev/src/scitex/db/_PostgreSQLMixins/_IndexMixin.py
+
+THIS_FILE = (
+    "/home/ywatanabe/proj/scitex_repo/src/scitex/db/_PostgreSQLMixins/_IndexMixin.py"
+)
+
+from typing import List
+import psycopg2
+
+
+class _IndexMixin:
+    def create_index(
+        self,
+        table_name: str,
+        column_names: List[str],
+        index_name: str = None,
+        unique: bool = False,
+    ) -> None:
+        try:
+            if index_name is None:
+                index_name = f"idx_{table_name}_{'_'.join(column_names)}"
+
+            unique_clause = "UNIQUE" if unique else ""
+            columns_str = ", ".join(column_names)
+
+            query = f"""
+            CREATE {unique_clause} INDEX IF NOT EXISTS {index_name}
+            ON {table_name} ({columns_str})
+            """
+            self.execute(query)
+
+        except (Exception, psycopg2.Error) as err:
+            raise ValueError(f"Failed to create index: {err}")
+
+    def drop_index(self, index_name: str) -> None:
+        try:
+            self.execute(f"DROP INDEX IF EXISTS {index_name}")
+        except (Exception, psycopg2.Error) as err:
+            raise ValueError(f"Failed to drop index: {err}")
+
+    def get_indexes(self, table_name: str = None) -> List[dict]:
+        try:
+            query = """
+            SELECT
+                schemaname,
+                tablename,
+                indexname,
+                indexdef
+            FROM
+                pg_indexes
+            """
+            if table_name:
+                query += f" WHERE tablename = '{table_name}'"
+
+            return self.execute(query).fetchall()
+
+        except (Exception, psycopg2.Error) as err:
+            raise ValueError(f"Failed to get indexes: {err}")
+
+
+# EOF