scitex 2.0.0__py2.py3-none-any.whl

scitex/db/_PostgreSQLMixins/_MaintenanceMixin.py

@@ -0,0 +1,175 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Timestamp: "2025-02-27 22:15:18 (ywatanabe)"
+# File: /home/ywatanabe/proj/scitex_dev/src/scitex/db/_PostgreSQLMixins/_MaintenanceMixin.py
+
+THIS_FILE = (
+    "/home/ywatanabe/proj/scitex_repo/src/scitex/db/_PostgreSQLMixins/_MaintenanceMixin.py"
+)
+
+import pandas as pd
+import contextlib
+from typing import ContextManager, Dict, List, Optional
+
+import psycopg2
+
+from .._BaseMixins._BaseMaintenanceMixin import _BaseMaintenanceMixin
+
+
+class _MaintenanceMixin(_BaseMaintenanceMixin):
+    @contextlib.contextmanager
+    def maintenance_lock(self) -> ContextManager[None]:
+        if not self._maintenance_lock.acquire(timeout=300):
+            raise TimeoutError("Could not acquire maintenance lock")
+        try:
+            yield
+        finally:
+            self._maintenance_lock.release()
+
+    def vacuum(self, table: Optional[str] = None, full: bool = False) -> None:
+        """Execute VACUUM on the specified table or entire database"""
+        with self.maintenance_lock():
+            try:
+                self._check_writable()
+                query = "VACUUM"
+                if full:
+                    query += " FULL"
+                if table:
+                    query += f" {table}"
+                self.execute(query)
+            except psycopg2.Error as err:
+                raise ValueError(f"Vacuum operation failed: {err}")
+
+    def analyze(self, table: Optional[str] = None) -> None:
+        """Update statistics for query optimization"""
+        with self.maintenance_lock():
+            try:
+                self._check_writable()
+                query = "ANALYZE"
+                if table:
+                    query += f" {table}"
+                self.execute(query)
+            except psycopg2.Error as err:
+                raise ValueError(f"Analyze operation failed: {err}")
+
+    def reindex(self, table: Optional[str] = None) -> None:
+        """Rebuild indexes"""
+        with self.maintenance_lock():
+            try:
+                self._check_writable()
+                if table:
+                    self.execute(f"REINDEX TABLE {table}")
+                else:
+                    self.execute("REINDEX DATABASE CURRENT_DATABASE()")
+            except psycopg2.Error as err:
+                raise ValueError(f"Reindex operation failed: {err}")
+
+    def get_table_size(self, table: str) -> str:
+        """Get the size of a specific table"""
+        try:
+            query = """
+                SELECT pg_size_pretty(pg_total_relation_size(%s))
+            """
+            self.execute(query, (table,))
+            return self.cursor.fetchone()[0]
+        except psycopg2.Error as err:
+            raise ValueError(f"Failed to get table size: {err}")
+
+    def get_database_size(self) -> str:
+        """Get the size of the current database"""
+        try:
+            query = """
+                SELECT pg_size_pretty(pg_database_size(current_database()))
+            """
+            self.execute(query)
+            return self.cursor.fetchone()[0]
+        except psycopg2.Error as err:
+            raise ValueError(f"Failed to get database size: {err}")
+
+    def get_table_info(self) -> List[Dict]:
+        """Get information about all tables in the database"""
+        try:
+            query = """
+                SELECT
+                    table_name,
+                    pg_size_pretty(pg_total_relation_size(quote_ident(table_name))) as size,
+                    (SELECT count(*) FROM information_schema.columns
+                     WHERE table_name=tables.table_name) as columns,
+                    (SELECT COUNT(*) FROM information_schema.table_constraints
+                     WHERE table_name=tables.table_name AND constraint_type='PRIMARY KEY') as has_pk
+                FROM information_schema.tables
+                WHERE table_schema = 'public'
+                ORDER BY table_name;
+            """
+            self.execute(query)
+            columns = [desc[0] for desc in self.cursor.description]
+            return [dict(zip(columns, row)) for row in self.cursor.fetchall()]
+        except psycopg2.Error as err:
+            raise ValueError(f"Failed to get table information: {err}")
+
+    def optimize(self, table: Optional[str] = None) -> None:
+        """Perform full optimization on table or database"""
+        with self.maintenance_lock():
+            try:
+                self.vacuum(table, full=True)
+                self.analyze(table)
+                self.reindex(table)
+            except ValueError as err:
+                raise ValueError(f"Optimization failed: {err}")
+
+    def get_summaries(
+        self,
+        table_names: Optional[List[str]] = None,
+        verbose: bool = True,
+        limit: int = 5,
+    ) -> Dict[str, pd.DataFrame]:
+
+        try:
+            if table_names is None:
+                table_names = self.get_table_names()
+            if isinstance(table_names, str):
+                table_names = [table_names]
+
+            sample_tables = {}
+            for table_name in table_names:
+                query = f"""
+                    SELECT *
+                    FROM {table_name}
+                    LIMIT {limit}
+                """
+                self.execute(query)
+                columns = [desc[0] for desc in self.cursor.description]
+                rows = self.cursor.fetchall()
+                table_sample = pd.DataFrame(rows, columns=columns)
+
+                for column in table_sample.columns:
+                    if table_sample[column].dtype == object:
+                        try:
+                            pd.to_datetime(
+                                table_sample[column], format="mixed", errors="raise"
+                            )
+                            continue
+                        except (ValueError, TypeError):
+                            pass
+
+                        # try:
+                        #     pd.to_datetime(table_sample[column])
+                        #     continue
+                        # except:
+                        #     pass
+
+                        if (
+                            table_sample[column]
+                            .apply(lambda x: isinstance(x, str))
+                            .all()
+                        ):
+                            continue
+
+                sample_tables[table_name] = table_sample
+
+            return sample_tables
+        except psycopg2.Error as err:
+            raise ValueError(f"Failed to get summaries: {err}")
+
+
+# EOF

scitex/db/_PostgreSQLMixins/_QueryMixin.py

@@ -0,0 +1,108 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Timestamp: "2025-02-27 22:15:24 (ywatanabe)"
+# File: /home/ywatanabe/proj/scitex_dev/src/scitex/db/_PostgreSQLMixins/_QueryMixin.py
+
+THIS_FILE = (
+    "/home/ywatanabe/proj/scitex_repo/src/scitex/db/_PostgreSQLMixins/_QueryMixin.py"
+)
+
+from typing import List, Dict, Any, Optional, Union, Tuple
+from .._BaseMixins._BaseQueryMixin import _BaseQueryMixin
+
+
+class _QueryMixin(_BaseQueryMixin):
+    def select(
+        self,
+        table: str,
+        columns: Optional[List[str]] = None,
+        where: Optional[str] = None,
+        params: Optional[tuple] = None,
+        order_by: Optional[str] = None,
+        limit: Optional[int] = None,
+    ) -> List[Dict[str, Any]]:
+        """Execute a SELECT query with optional conditions"""
+        cols_str = "*" if not columns else ", ".join(columns)
+        query = f"SELECT {cols_str} FROM {table}"
+
+        if where:
+            query += f" WHERE {where}"
+        if order_by:
+            query += f" ORDER BY {order_by}"
+        if limit:
+            query += f" LIMIT {limit}"
+
+        self.execute(query, params)
+        columns = [desc[0] for desc in self.cursor.description]
+        return [dict(zip(columns, row)) for row in self.cursor.fetchall()]
+
+    def insert(self, table: str, data: Dict[str, Any]) -> None:
+        """Insert a single record into a table"""
+        self._check_writable()
+        columns = list(data.keys())
+        values = list(data.values())
+        placeholders = ["%s"] * len(values)
+
+        query = f"""
+            INSERT INTO {table}
+            ({', '.join(columns)})
+            VALUES ({', '.join(placeholders)})
+        """
+
+        self.execute(query, tuple(values))
+
+    def update(
+        self,
+        table: str,
+        data: Dict[str, Any],
+        where: str,
+        params: Optional[tuple] = None,
+    ) -> int:
+        """Update records in a table"""
+        self._check_writable()
+        set_items = [f"{k} = %s" for k in data.keys()]
+        values = list(data.values())
+
+        query = f"""
+            UPDATE {table}
+            SET {', '.join(set_items)}
+            WHERE {where}
+        """
+
+        if params:
+            values.extend(params)
+
+        self.execute(query, tuple(values))
+        return self.cursor.rowcount
+
+    def delete(self, table: str, where: str, params: Optional[tuple] = None) -> int:
+        """Delete records from a table"""
+        self._check_writable()
+        query = f"DELETE FROM {table} WHERE {where}"
+        self.execute(query, params)
+        return self.cursor.rowcount
+
+    def execute_query(
+        self, query: str, params: Optional[tuple] = None
+    ) -> List[Dict[str, Any]]:
+        """Execute a custom query and return results as dictionaries"""
+        self.execute(query, params)
+
+        if self.cursor.description:  # If the query returns results
+            columns = [desc[0] for desc in self.cursor.description]
+            return [dict(zip(columns, row)) for row in self.cursor.fetchall()]
+        return []
+
+    def count(
+        self, table: str, where: Optional[str] = None, params: Optional[tuple] = None
+    ) -> int:
+        """Count records in a table"""
+        query = f"SELECT COUNT(*) FROM {table}"
+        if where:
+            query += f" WHERE {where}"
+
+        self.execute(query, params)
+        return self.cursor.fetchone()[0]
+
+
+# EOF

scitex/db/_PostgreSQLMixins/_RowMixin.py

@@ -0,0 +1,75 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Timestamp: "2025-02-27 22:15:30 (ywatanabe)"
+# File: /home/ywatanabe/proj/scitex_dev/src/scitex/db/_PostgreSQLMixins/_RowMixin.py
+
+THIS_FILE = "/home/ywatanabe/proj/scitex_repo/src/scitex/db/_PostgreSQLMixins/_RowMixin.py"
+
+from typing import List, Optional
+import pandas as pd
+import psycopg2
+
+
+class _RowMixin:
+    def get_rows(
+        self,
+        table_name: str,
+        columns: List[str] = None,
+        where: str = None,
+        order_by: str = None,
+        limit: Optional[int] = None,
+        offset: Optional[int] = None,
+        return_as: str = "dataframe",
+    ):
+        try:
+            if columns is None:
+                columns_str = "*"
+            elif isinstance(columns, str):
+                columns_str = f'"{columns}"'
+            else:
+                columns_str = ", ".join(f'"{col}"' for col in columns)
+
+            query_parts = [f"SELECT {columns_str} FROM {table_name}"]
+
+            if where:
+                query_parts.append(f"WHERE {where}")
+            if order_by:
+                query_parts.append(f"ORDER BY {order_by}")
+            if limit is not None:
+                query_parts.append(f"LIMIT {limit}")
+            if offset is not None:
+                query_parts.append(f"OFFSET {offset}")
+
+            query = " ".join(query_parts)
+            self.cursor.execute(query)
+
+            column_names = [desc[0] for desc in self.cursor.description]
+            data = self.cursor.fetchall()
+
+            if return_as == "list":
+                return data
+            elif return_as == "dict":
+                return [dict(zip(column_names, row)) for row in data]
+            else:
+                return pd.DataFrame(data, columns=column_names)
+
+        except (Exception, psycopg2.Error) as err:
+            raise ValueError(f"Query execution failed: {err}")
+
+    def get_row_count(self, table_name: str = None, where: str = None) -> int:
+        try:
+            if table_name is None:
+                raise ValueError("Table name must be specified")
+
+            query = f"SELECT COUNT(*) FROM {table_name}"
+            if where:
+                query += f" WHERE {where}"
+
+            self.cursor.execute(query)
+            return self.cursor.fetchone()[0]
+
+        except (Exception, psycopg2.Error) as err:
+            raise ValueError(f"Failed to get row count: {err}")
+
+
+# EOF

scitex/db/_PostgreSQLMixins/_SchemaMixin.py

@@ -0,0 +1,126 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Timestamp: "2025-02-27 22:14:23 (ywatanabe)"
+# File: /home/ywatanabe/proj/scitex_dev/src/scitex/db/_PostgreSQLMixins/_SchemaMixin.py
+
+THIS_FILE = (
+    "/home/ywatanabe/proj/scitex_repo/src/scitex/db/_PostgreSQLMixins/_SchemaMixin.py"
+)
+
+from typing import List, Dict, Any, Optional
+from .._BaseMixins._BaseSchemaMixin import _BaseSchemaMixin
+
+
+class _SchemaMixin(_BaseSchemaMixin):
+    def get_tables(self) -> List[str]:
+        """Get all tables in the current database"""
+        query = """
+            SELECT table_name
+            FROM information_schema.tables
+            WHERE table_schema = 'public'
+        """
+        self.execute(query)
+        return [row[0] for row in self.cursor.fetchall()]
+
+    def get_columns(self, table: str) -> List[Dict[str, Any]]:
+        """Get detailed information about columns in a table"""
+        query = """
+            SELECT column_name, data_type, is_nullable, column_default
+            FROM information_schema.columns
+            WHERE table_schema = 'public' AND table_name = %s
+            ORDER BY ordinal_position
+        """
+        self.execute(query, (table,))
+        columns = [desc[0] for desc in self.cursor.description]
+        return [dict(zip(columns, row)) for row in self.cursor.fetchall()]
+
+    def get_primary_keys(self, table: str) -> List[str]:
+        """Get primary key columns for a table"""
+        query = """
+            SELECT a.attname
+            FROM pg_index i
+            JOIN pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = ANY(i.indkey)
+            WHERE i.indrelid = %s::regclass AND i.indisprimary
+        """
+        self.execute(query, (table,))
+        return [row[0] for row in self.cursor.fetchall()]
+
+    def get_foreign_keys(self, table: str) -> List[Dict[str, Any]]:
+        """Get foreign key constraints for a table"""
+        query = """
+            SELECT
+                kcu.column_name,
+                ccu.table_name AS foreign_table_name,
+                ccu.column_name AS foreign_column_name
+            FROM information_schema.table_constraints AS tc
+            JOIN information_schema.key_column_usage AS kcu
+                ON tc.constraint_name = kcu.constraint_name
+            JOIN information_schema.constraint_column_usage AS ccu
+                ON ccu.constraint_name = tc.constraint_name
+            WHERE tc.constraint_type = 'FOREIGN KEY'
+                AND tc.table_name = %s
+        """
+        self.execute(query, (table,))
+        columns = [desc[0] for desc in self.cursor.description]
+        return [dict(zip(columns, row)) for row in self.cursor.fetchall()]
+
+    def get_indexes(self, table: str) -> List[Dict[str, Any]]:
+        """Get all indexes for a table"""
+        query = """
+            SELECT
+                i.relname as index_name,
+                a.attname as column_name,
+                ix.indisunique as is_unique,
+                ix.indisprimary as is_primary
+            FROM pg_class t
+            JOIN pg_index ix ON t.oid = ix.indrelid
+            JOIN pg_class i ON i.oid = ix.indexrelid
+            JOIN pg_attribute a ON a.attrelid = t.oid AND a.attnum = ANY(ix.indkey)
+            WHERE t.relname = %s
+        """
+        self.execute(query, (table,))
+        columns = [desc[0] for desc in self.cursor.description]
+        return [dict(zip(columns, row)) for row in self.cursor.fetchall()]
+
+    def table_exists(self, table: str) -> bool:
+        """Check if a table exists"""
+        query = """
+            SELECT EXISTS (
+                SELECT FROM information_schema.tables
+                WHERE table_schema = 'public'
+                AND table_name = %s
+            )
+        """
+        self.execute(query, (table,))
+        return self.cursor.fetchone()[0]
+
+    def column_exists(self, table: str, column: str) -> bool:
+        """Check if a column exists in a table"""
+        query = """
+            SELECT EXISTS (
+                SELECT FROM information_schema.columns
+                WHERE table_schema = 'public'
+                AND table_name = %s
+                AND column_name = %s
+            )
+        """
+        self.execute(query, (table, column))
+        return self.cursor.fetchone()[0]
+
+    def create_index(
+        self, table: str, columns: List[str], index_name: Optional[str] = None
+    ) -> None:
+        """Create an index on specified columns"""
+        self._check_writable()
+        if not index_name:
+            index_name = f"idx_{table}_{'_'.join(columns)}"
+        columns_str = ", ".join(columns)
+        self.execute(f"CREATE INDEX {index_name} ON {table} ({columns_str})")
+
+    def drop_index(self, index_name: str) -> None:
+        """Drop an index"""
+        self._check_writable()
+        self.execute(f"DROP INDEX IF EXISTS {index_name}")
+
+
+# EOF

scitex/db/_PostgreSQLMixins/_TableMixin.py

@@ -0,0 +1,176 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Timestamp: "2025-02-27 22:15:38 (ywatanabe)"
+# File: /home/ywatanabe/proj/scitex_dev/src/scitex/db/_PostgreSQLMixins/_TableMixin.py
+
+THIS_FILE = (
+    "/home/ywatanabe/proj/scitex_repo/src/scitex/db/_PostgreSQLMixins/_TableMixin.py"
+)
+
+from typing import Any, Dict, List, Union
+import psycopg2
+
+
+class _TableMixin:
+    def create_table(
+        self,
+        table_name: str,
+        columns: Dict[str, str],
+        foreign_keys: List[Dict[str, str]] = None,
+        if_not_exists: bool = True,
+    ) -> None:
+        try:
+            exists_clause = "IF NOT EXISTS" if if_not_exists else ""
+            columns_def = [
+                f"{col_name} {col_type}" for col_name, col_type in columns.items()
+            ]
+
+            if foreign_keys:
+                for fk in foreign_keys:
+                    columns_def.append(
+                        f"FOREIGN KEY ({fk['column']}) REFERENCES {fk['references']}({fk['referenced_column']})"
+                    )
+
+            columns_str = ", ".join(columns_def)
+            query = f"CREATE TABLE {exists_clause} {table_name} ({columns_str})"
+            self.execute(query)
+
+        except (Exception, psycopg2.Error) as err:
+            raise ValueError(f"Failed to create table: {err}")
+
+    def drop_table(self, table_name: str, if_exists: bool = True) -> None:
+        try:
+            exists_clause = "IF EXISTS" if if_exists else ""
+            self.execute(f"DROP TABLE {exists_clause} {table_name}")
+
+        except (Exception, psycopg2.Error) as err:
+            raise ValueError(f"Failed to drop table: {err}")
+
+    def rename_table(self, old_name: str, new_name: str) -> None:
+        try:
+            self.execute(f"ALTER TABLE {old_name} RENAME TO {new_name}")
+
+        except (Exception, psycopg2.Error) as err:
+            raise ValueError(f"Failed to rename table: {err}")
+
+    def add_columns(
+        self,
+        table_name: str,
+        columns: Dict[str, str],
+        default_values: Dict[str, Any] = None,
+    ) -> None:
+        try:
+            for col_name, col_type in columns.items():
+                default_value = default_values.get(col_name) if default_values else None
+                default_clause = (
+                    f" DEFAULT {default_value}" if default_value is not None else ""
+                )
+
+                self.execute(
+                    f"ALTER TABLE {table_name} ADD COLUMN {col_name} {col_type}{default_clause}"
+                )
+
+        except (Exception, psycopg2.Error) as err:
+            raise ValueError(f"Failed to add columns: {err}")
+
+    def add_column(
+        self,
+        table_name: str,
+        column_name: str,
+        column_type: str,
+        default_value: Any = None,
+    ) -> None:
+        try:
+            default_clause = (
+                f" DEFAULT {default_value}" if default_value is not None else ""
+            )
+            self.execute(
+                f"ALTER TABLE {table_name} ADD COLUMN {column_name} {column_type}{default_clause}"
+            )
+
+        except (Exception, psycopg2.Error) as err:
+            raise ValueError(f"Failed to add column: {err}")
+
+    def drop_columns(
+        self, table_name: str, columns: Union[str, List[str]], if_exists: bool = True
+    ) -> None:
+        try:
+            if isinstance(columns, str):
+                columns = [columns]
+
+            exists_clause = "IF EXISTS" if if_exists else ""
+            for column in columns:
+                self.execute(
+                    f"ALTER TABLE {table_name} DROP COLUMN {exists_clause} {column}"
+                )
+
+        except (Exception, psycopg2.Error) as err:
+            raise ValueError(f"Failed to drop columns: {err}")
+
+    def get_table_names(self) -> List[str]:
+        try:
+            query = """
+            SELECT table_name
+            FROM information_schema.tables
+            WHERE table_schema = 'public'
+            """
+            self.cursor.execute(query)
+            return [row[0] for row in self.cursor.fetchall()]
+
+        except (Exception, psycopg2.Error) as err:
+            raise ValueError(f"Failed to get table names: {err}")
+
+    def get_table_schema(self, table_name: str):
+        try:
+            query = f"""
+            SELECT column_name, data_type, is_nullable
+            FROM information_schema.columns
+            WHERE table_name = '{table_name}'
+            """
+            return self.execute(query).fetchall()
+
+        except (Exception, psycopg2.Error) as err:
+            raise ValueError(f"Failed to get table schema: {err}")
+
+    def get_primary_key(self, table_name: str) -> str:
+        try:
+            query = f"""
+            SELECT a.attname
+            FROM   pg_index i
+            JOIN   pg_attribute a ON a.attrelid = i.indrelid
+                                AND a.attnum = ANY(i.indkey)
+            WHERE  i.indrelid = '{table_name}'::regclass
+            AND    i.indisprimary
+            """
+            result = self.execute(query).fetchone()
+            return result[0] if result else None
+
+        except (Exception, psycopg2.Error) as err:
+            raise ValueError(f"Failed to get primary key: {err}")
+
+    def get_table_stats(self, table_name: str) -> Dict[str, int]:
+        try:
+            stats = {}
+            stats["row_count"] = self.get_row_count(table_name)
+
+            size_query = f"""
+            SELECT pg_total_relation_size('{table_name}') as total_size,
+                   pg_table_size('{table_name}') as table_size,
+                   pg_indexes_size('{table_name}') as index_size
+            """
+            size_result = self.execute(size_query).fetchone()
+
+            stats.update(
+                {
+                    "total_size": size_result[0],
+                    "table_size": size_result[1],
+                    "index_size": size_result[2],
+                }
+            )
+            return stats
+
+        except (Exception, psycopg2.Error) as err:
+            raise ValueError(f"Failed to get table stats: {err}")
+
+
+# EOF