thoth-dbmanager 0.4.0__py3-none-any.whl → 0.4.1__py3-none-any.whl

thoth_dbmanager/helpers/schema.py

@@ -0,0 +1,376 @@
+import logging
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Tuple
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ColumnInfo:
+    """
+    Represents metadata for a single column in a database table.
+
+    Attributes:
+        original_column_name (str): The original name of the column.
+        column_name (str): The standardized name of the column.
+        column_description (str): A description of the column.
+        data_format (str): The format of the data in the column.
+        value_description (str): A description of the values in the column.
+        type (str): The data type of the column.
+        examples (List[str]): Example values from the column.
+        primary_key (bool): Whether the column is a primary key.
+        foreign_keys (List[Tuple[str, str]]): Foreign keys referencing other tables and columns.
+        referenced_by (List[Tuple[str, str]]): Columns in other tables that reference this column.
+    """
+
+    original_column_name: str = ""
+    column_name: str = ""
+    column_description: str = ""
+    generated_comment: str = ""
+    value_description: str = ""
+    data_format: str = ""
+    type: str = ""
+    examples: List[str] = field(default_factory=list)
+    primary_key: bool = False
+    foreign_keys: List[Tuple[str, str]] = field(default_factory=list)
+    referenced_by: List[Tuple[str, str]] = field(default_factory=list)
+
+
+def set_field(column_info: ColumnInfo, field_name: str, value: Any) -> None:
+    """
+    Sets a field in the ColumnInfo dataclass.
+
+    Args:
+        column_info (ColumnInfo): The ColumnInfo instance to update.
+        field_name (str): The field name to set.
+        value (Any): The value to set for the field.
+
+    Raises:
+        ValueError: If the field_name is not a valid field of ColumnInfo.
+    """
+    if field_name in column_info.__dataclass_fields__:
+        setattr(column_info, field_name, value)
+    else:
+        raise ValueError(f"{field_name} is not a valid field of ColumnInfo")
+
+
+@dataclass
+class TableSchema:
+    """
+    Represents the schema of a single table in a database.
+
+    Attributes:
+        columns (Dict[str, ColumnInfo]): A dictionary mapping column names to their metadata.
+    """
+
+    columns: Dict[str, ColumnInfo] = field(default_factory=dict)
+
+
+def get_primary_keys(table_schema: TableSchema) -> List[str]:
+    """
+    Retrieves the primary key columns from a table schema.
+
+    Args:
+        table_schema (TableSchema): The table schema to analyze.
+
+    Returns:
+        List[str]: A list of primary key column names.
+    """
+    return [name for name, info in table_schema.columns.items() if info.primary_key]
+
+
+@dataclass
+class DatabaseSchema:
+    """
+    Represents the schema of an entire database, consisting of multiple tables.
+
+    Attributes:
+        tables (Dict[str, TableSchema]): A dictionary mapping table names to their schemas.
+    """
+
+    tables: Dict[str, TableSchema] = field(default_factory=dict)
+
+    @classmethod
+    def from_table_names(cls, table_names: List[str]) -> "DatabaseSchema":
+        """
+        Creates a DatabaseSchema from a list of table names.
+
+        Args:
+            table_names (List[str]): The names of the tables to include in the schema.
+
+        Returns:
+            DatabaseSchema: The constructed database schema.
+        """
+        return cls(tables={name: TableSchema() for name in table_names})
+
+    @classmethod
+    def from_schema_dict(
+            cls, schema_dict: Dict[str, List[str]]
+    ) -> "DatabaseSchema":
+        """
+        Creates a DatabaseSchema from a dictionary mapping table names to lists of column names.
+
+        Args:
+            schema_dict (Dict[str, List[str]]): The schema dictionary to convert.
+
+        Returns:
+            DatabaseSchema: The constructed database schema.
+        """
+        return cls(
+            tables={
+                table_name: TableSchema(
+                    columns={column_name: ColumnInfo() for column_name in column_names}
+                )
+                for table_name, column_names in schema_dict.items()
+            }
+        )
+
+    @classmethod
+    def from_schema_dict_with_examples(
+        cls, schema_dict_with_info: Dict[str, Dict[str, List[str]]]
+    ) -> "DatabaseSchema":
+        """
+        Creates a DatabaseSchema from a dictionary with example values for each column.
+
+        Args:
+            schema_dict_with_info (Dict[str, Dict[str, List[str]]]): The schema dictionary with example values.
+
+        Returns:
+            DatabaseSchema: The constructed database schema.
+        """
+        return cls(
+            tables={
+                table_name: TableSchema(
+                    columns={
+                        column_name: ColumnInfo(examples=column_info)
+                        for column_name, column_info in column_dict.items()
+                    }
+                )
+                for table_name, column_dict in schema_dict_with_info.items()
+            }
+        )
+
+    @classmethod
+    def from_schema_dict_with_descriptions(
+        cls,
+        tentative_schema: Dict[str, Dict[str, Dict[str, Any]]],
+        schema_dict_with_info: Dict[str, Dict[str, Dict[str, Any]]]
+    ) -> "DatabaseSchema":
+        """
+        Creates a DatabaseSchema from a dictionary with detailed information for each column.
+
+        Args:
+            tentative_schema (Dict[str, Dict[str, Dict[str, Any]]]): The base schema structure
+            schema_dict_with_info (Dict[str, Dict[str, Dict[str, Any]]]): The schema dictionary with detailed information
+
+        Returns:
+            DatabaseSchema: The constructed database schema
+        """
+        # Convert new schema format to simplified format for initial creation
+        simplified_schema = {
+            table_name: list(table_info["columns"].keys())
+            for table_name, table_info in tentative_schema.items()
+        }
+        
+        database_schema = cls.from_schema_dict(simplified_schema)
+        
+        for table_name, table_info in schema_dict_with_info.items():
+            if table_name in database_schema.tables:
+                for column_name, info in table_info.items():
+                    if column_name in database_schema.tables[table_name].columns:
+                        column_info = database_schema.tables[table_name].columns[column_name]
+                        for field_name, value in info.items():
+                            set_field(column_info, field_name, value)
+        
+        return database_schema
+
+    def get_actual_table_name(
+            self, table_name: str
+    ) -> Optional[str]:
+        """
+        Retrieves the actual table name matching the provided name, case-insensitive.
+
+        Args:
+            table_name (str): The name of the table to search for.
+
+        Returns:
+            Optional[str]: The actual table name if found, otherwise None.
+        """
+        table_name_lower = table_name.lower()
+        return next(
+            (name for name in self.tables if name.lower() == table_name_lower), None
+        )
+
+    def get_table_info(
+            self, table_name: str
+    ) -> Optional[TableSchema]:
+        """
+        Retrieves the TableSchema object for the specified table name.
+
+        Args:
+            table_name (str): The name of the table to retrieve.
+
+        Returns:
+            Optional[TableSchema]: The TableSchema if found, otherwise None.
+        """
+        actual_name = self.get_actual_table_name(table_name)
+        return self.tables.get(actual_name)
+
+    def get_actual_column_name(
+        self, table_name: str, column_name: str
+    ) -> Optional[str]:
+        """
+        Retrieves the actual column name matching the provided name, case-insensitive.
+
+        Args:
+            table_name (str): The name of the table containing the column.
+            column_name (str): The name of the column to search for.
+
+        Returns:
+            Optional[str]: The actual column name if found, otherwise None.
+        """
+        table_info = self.get_table_info(table_name)
+        if table_info:
+            column_name_lower = column_name.lower()
+            return next(
+                (
+                    name
+                    for name in table_info.columns
+                    if name.lower() == column_name_lower
+                ),
+                None,
+            )
+        return None
+
+    def get_column_info(
+        self, table_name: str, column_name: str
+    ) -> Optional[ColumnInfo]:
+        """
+        Retrieves the ColumnInfo object for the specified column in a table.
+
+        Args:
+            table_name (str): The name of the table containing the column.
+            column_name (str): The name of the column to retrieve.
+
+        Returns:
+            Optional[ColumnInfo]: The ColumnInfo if found, otherwise None.
+        """
+        actual_name = self.get_actual_column_name(table_name, column_name)
+        if actual_name:
+            return self.tables[table_name].columns[actual_name]
+        return None
+
+    def set_columns_info(
+        self, schema_with_info: Dict[str, Dict[str, Dict[str, Any]]]
+    ) -> None:
+        """
+        Sets detailed information for columns in the schema.
+
+        Args:
+            schema_with_info (Dict[str, Dict[str, Dict[str, Any]]]): The schema information to set.
+        """
+        for table_name, columns_info in schema_with_info.items():
+            table_info = self.get_table_info(table_name)
+            if table_info is None:
+                logger.warning(f"Table {table_name} not found in the schema")
+                continue
+            for column_name, info in columns_info.items():
+                actual_name = self.get_actual_column_name(table_name, column_name)
+                if actual_name is None:
+                    logger.warning(f"Column {column_name} not found in table {table_name}")
+                    continue
+                schema_column_info = table_info.columns[actual_name]
+                for field_name, value in info.items():
+                    set_field(schema_column_info, field_name, value)
+
+    def subselect_schema(
+        self, selected_database_schema: "DatabaseSchema"
+    ) -> "DatabaseSchema":
+        """
+        Creates a new DatabaseSchema containing only the selected tables and columns.
+
+        Args:
+            selected_database_schema (DatabaseSchema): The schema to subselect from.
+
+        Returns:
+            DatabaseSchema: The new subselected database schema.
+        """
+        new_schema = DatabaseSchema({})
+        for table_name, table_info in selected_database_schema.tables.items():
+            actual_table_name = self.get_actual_table_name(table_name)
+            if actual_table_name is None:
+                logger.warning(f"Table {table_name} not found in the schema")
+                continue
+            new_table_info = TableSchema()
+            for column_name, column_info in table_info.columns.items():
+                actual_column_name = self.get_actual_column_name(
+                    table_name, column_name
+                )
+                if actual_column_name is None:
+                    logger.warning(f"Column {column_name} not found in table {table_name}")
+                    continue
+                new_table_info.columns[actual_column_name] = column_info
+            new_schema.tables[actual_table_name] = new_table_info
+        return new_schema
+
+    def add_info_from_schema(
+        self, schema: "DatabaseSchema", field_names: List[str]
+    ) -> None:
+        """
+        Adds additional field information from another schema to the current schema.
+
+        Args:
+            schema (DatabaseSchema): The schema to copy information from.
+            field_names (List[str]): The list of field names to copy.
+        """
+        for table_name, table_info in self.tables.items():
+            actual_table_name = schema.get_actual_table_name(table_name)
+            if actual_table_name is None:
+                continue
+            for column_name, column_info in table_info.columns.items():
+                actual_column_name = schema.get_actual_column_name(
+                    table_name, column_name
+                )
+                if actual_column_name is None:
+                    continue
+                new_column_info = schema.tables[actual_table_name].columns[
+                    actual_column_name
+                ]
+                for field_name in field_names:
+                    set_field(
+                        column_info, field_name, getattr(new_column_info, field_name)
+                    )
+
+    def to_dict(self) -> Dict[str, List[str]]:
+        """
+        Converts the DatabaseSchema to a dictionary representation.
+
+        Returns:
+            Dict[str, List[str]]: The dictionary representation of the schema.
+        """
+        return {
+            table_name: list(table_info.columns.keys())
+            for table_name, table_info in self.tables.items()
+        }
+
+    def validate_schema(self) -> List[str]:
+        """
+        Validates the schema for integrity and returns a list of validation errors.
+        
+        Returns:
+            List[str]: A list of validation error messages, empty if no errors.
+        """
+        errors = []
+        
+        # Check for foreign key references to non-existent tables/columns
+        for table_name, table_info in self.tables.items():
+            for column_name, column_info in table_info.columns.items():
+                for fk_table, fk_column in column_info.foreign_keys:
+                    if not self.get_actual_table_name(fk_table):
+                        errors.append(f"Foreign key in {table_name}.{column_name} references non-existent table {fk_table}")
+                        continue
+                        
+                    if not self.get_actual_column_name(fk_table, fk_column):
+                        errors.append(f"Foreign key in {table_name}.{column_name} references non-existent column {fk_table}.{fk_column}")
+        
+        return errors

thoth_dbmanager/helpers/search.py

@@ -0,0 +1,117 @@
+import logging
+import pickle
+from pathlib import Path
+from typing import Dict, List, Tuple
+
+from datasketch import MinHash, MinHashLSH
+
+from .preprocess_values import _create_minhash
+
+
+### Database value similarity ###
+
+
+def _jaccard_similarity(m1: MinHash, m2: MinHash) -> float:
+    """
+    Computes the Jaccard similarity between two MinHash objects.
+
+    Args:
+        m1 (MinHash): The first MinHash object.
+        m2 (MinHash): The second MinHash object.
+
+    Returns:
+        float: The Jaccard similarity between the two MinHash objects.
+    """
+    return m1.jaccard(m2)
+
+
+def load_db_lsh(
+    db_directory_path: str,
+) -> Tuple[MinHashLSH, Dict[str, Tuple[MinHash, str, str, str]]]:
+    """
+    Loads the LSH and MinHashes from the preprocessed files in the specified directory.
+    
+    This function maintains backward compatibility while potentially using the new LSH manager.
+
+    Args:
+        db_directory_path (str): The path to the database directory.
+
+    Returns:
+        Tuple[MinHashLSH, Dict[str, Tuple[MinHash, str, str, str]]]: The LSH object and the dictionary of MinHashes.
+
+    Raises:
+        Exception: If there is an error loading the LSH or MinHashes.
+    """
+    db_id = Path(db_directory_path).name
+    
+    try:
+        # Try using the new LSH manager first
+        from ..lsh.manager import LshManager
+        lsh_manager = LshManager(Path(db_directory_path))
+        if lsh_manager.load_lsh():
+            return lsh_manager.lsh, lsh_manager.minhashes
+    except Exception as e:
+        logging.warning(f"New LSH manager failed, falling back to old method: {e}")
+    
+    # Fallback to old method
+    try:
+        with open(
+            Path(db_directory_path) / "preprocessed" / f"{db_id}_lsh.pkl", "rb"
+        ) as file:
+            lsh = pickle.load(file)
+        with open(
+            Path(db_directory_path) / "preprocessed" / f"{db_id}_minhashes.pkl", "rb"
+        ) as file:
+            minhashes = pickle.load(file)
+        return lsh, minhashes
+    except Exception as e:
+        logging.error(f"Error loading LSH for {db_id}: {e}")
+        raise e
+
+
+def _query_lsh(
+    lsh: MinHashLSH,
+    minhashes: Dict[str, Tuple[MinHash, str, str, str]],
+    keyword: str,
+    signature_size: int = 30,
+    n_gram: int = 3,
+    top_n: int = 10,
+) -> Dict[str, Dict[str, List[str]]]:
+    """
+    Queries the LSH for similar values to the given keyword and returns the top results.
+
+    Args:
+        lsh (MinHashLSH): The LSH object.
+        minhashes (Dict[str, Tuple[MinHash, str, str, str]]): The dictionary of MinHashes.
+        keyword (str): The keyword to search for.
+        signature_size (int, optional): The size of the MinHash signature.
+        n_gram (int, optional): The n-gram size for the MinHash.
+        top_n (int, optional): The number of top results to return.
+
+    Returns:
+        Dict[str, Dict[str, List[str]]]: A dictionary containing the top similar values.
+        Example:{
+        'table_name1': {
+            'column_name1': ['value1', 'value2', 'value3'],
+            'column_name2': ['value4', 'value5']
+        },
+        'table_name2': {
+    """
+    query_minhash = _create_minhash(signature_size, keyword, n_gram)
+    results = lsh.query(query_minhash)
+    similarities = [
+        (result, _jaccard_similarity(query_minhash, minhashes[result][0]))
+        for result in results
+    ]
+    similarities = sorted(similarities, key=lambda x: x[1], reverse=True)[:top_n]
+
+    similar_values_trimmed: Dict[str, Dict[str, List[str]]] = {}
+    for result, similarity in similarities:
+        table_name, column_name, value = minhashes[result][1:] #type: ignore
+        if table_name not in similar_values_trimmed:
+            similar_values_trimmed[table_name] = {}
+        if column_name not in similar_values_trimmed[table_name]:
+            similar_values_trimmed[table_name][column_name] = []
+        similar_values_trimmed[table_name][column_name].append(value)
+
+    return similar_values_trimmed

thoth_dbmanager/lsh/__init__.py

@@ -0,0 +1,21 @@
+"""
+LSH (Locality Sensitive Hashing) module for database-independent LSH management.
+"""
+
+from .storage import LshStorageStrategy, PickleStorage
+from .manager import LshManager
+from .factory import LshFactory, make_db_lsh
+from .core import create_minhash, skip_column, jaccard_similarity, create_lsh_index, query_lsh_index
+
+__all__ = [
+    "LshStorageStrategy",
+    "PickleStorage", 
+    "LshManager",
+    "LshFactory",
+    "make_db_lsh",
+    "create_minhash",
+    "skip_column", 
+    "jaccard_similarity",
+    "create_lsh_index",
+    "query_lsh_index"
+]

thoth_dbmanager/lsh/core.py

@@ -0,0 +1,182 @@
+"""
+Core LSH functionality extracted from helpers.
+"""
+
+import logging
+from typing import Dict, List, Tuple
+
+from datasketch import MinHash, MinHashLSH
+from tqdm import tqdm
+
+
+def create_minhash(signature_size: int, string: str, n_gram: int) -> MinHash:
+    """
+    Creates a MinHash object for a given string.
+
+    Args:
+        signature_size (int): The size of the MinHash signature.
+        string (str): The input string to create the MinHash for.
+        n_gram (int): The n-gram size for the MinHash.
+
+    Returns:
+        MinHash: The MinHash object for the input string.
+    """
+    m = MinHash(num_perm=signature_size)
+    for d in [string[i : i + n_gram] for i in range(len(string) - n_gram + 1)]:
+        m.update(d.encode("utf8"))
+    return m
+
+
+def skip_column(column_name: str, column_values: List[str]) -> bool:
+    """
+    Determines whether to skip processing a column based on its values.
+
+    Args:
+        column_name (str): The name of the column.
+        column_values (List[str]): The list of values in the column.
+
+    Returns:
+        bool: True if the column should be skipped, False otherwise.
+    """
+    if "name" in column_name.lower():
+        return False
+    sum_of_lengths = sum(len(value) for value in column_values)
+    average_length = sum_of_lengths / len(column_values)
+    return (sum_of_lengths > 50000) and (average_length > 20)
+
+
+def jaccard_similarity(m1: MinHash, m2: MinHash) -> float:
+    """
+    Computes the Jaccard similarity between two MinHash objects.
+
+    Args:
+        m1 (MinHash): The first MinHash object.
+        m2 (MinHash): The second MinHash object.
+
+    Returns:
+        float: The Jaccard similarity between the two MinHash objects.
+    """
+    return m1.jaccard(m2)
+
+
+def create_lsh_index(
+    unique_values: Dict[str, Dict[str, List[str]]],
+    signature_size: int,
+    n_gram: int,
+    threshold: float,
+    verbose: bool = True,
+) -> Tuple[MinHashLSH, Dict[str, Tuple[MinHash, str, str, str]]]:
+    """
+    Creates a MinHash Locality-Sensitive Hashing (LSH) index from unique values in a database.
+
+    This function processes unique values from database tables and columns, creates MinHash
+    signatures for each value, and builds an LSH index for efficient similarity search.
+
+    Args:
+        unique_values (Dict[str, Dict[str, List[str]]]): A nested dictionary containing unique values
+            from the database. The structure is {table_name: {column_name: [values]}}.
+        signature_size (int): The number of permutations to use in the MinHash signatures.
+        n_gram (int): The size of n-grams to use when creating MinHash signatures.
+        threshold (float): The similarity threshold for the LSH index. Values closer to 1 require
+            higher similarity for matches.
+        verbose (bool, optional): If True, displays a progress bar during processing. Defaults to True.
+
+    Returns:
+        Tuple[MinHashLSH, Dict[str, Tuple[MinHash, str, str, str]]]: A tuple containing:
+            - MinHashLSH: The constructed LSH index.
+            - Dict[str, Tuple[MinHash, str, str, str]]: A dictionary mapping unique keys to tuples
+              containing (MinHash object, table name, column name, original value).
+
+    Raises:
+        Exception: If an error occurs during LSH creation, it's logged but not raised.
+
+    Note:
+        This function uses the datasketch library for MinHash and LSH operations.
+    """
+    lsh = MinHashLSH(threshold=threshold, num_perm=signature_size)
+    minhashes: Dict[str, Tuple[MinHash, str, str, str]] = {}
+    try:
+        total_unique_values = sum(
+            len(column_values)
+            for table_values in unique_values.values()
+            for column_values in table_values.values()
+        )
+        logging.info(f"Total unique values: {total_unique_values}")
+
+        progress_bar = (
+            tqdm(total=total_unique_values, desc="Creating LSH") if verbose else None
+        )
+
+        for table_name, table_values in unique_values.items():
+            for column_name, column_values in table_values.items():
+                if column_name.lower() == "doctype":
+                    print("=" * 20)
+                    print("Doctype found")
+                    print("=" * 20)
+                logging.info(
+                    f"Processing {table_name} - {column_name} - {len(column_values)}"
+                )
+
+                for id, value in enumerate(column_values):
+                    minhash = create_minhash(signature_size, value, n_gram)
+                    minhash_key = f"{table_name}_{column_name}_{id}"
+                    minhashes[minhash_key] = (minhash, table_name, column_name, value)
+                    lsh.insert(minhash_key, minhash)
+
+                    if verbose:
+                        progress_bar.update(1)
+
+        if verbose:
+            progress_bar.close()
+    except Exception as e:
+        logging.error(f"Error creating LSH: {e}")
+
+    return lsh, minhashes
+
+
+def query_lsh_index(
+    lsh: MinHashLSH,
+    minhashes: Dict[str, Tuple[MinHash, str, str, str]],
+    keyword: str,
+    signature_size: int = 30,
+    n_gram: int = 3,
+    top_n: int = 10,
+) -> Dict[str, Dict[str, List[str]]]:
+    """
+    Queries the LSH for similar values to the given keyword and returns the top results.
+
+    Args:
+        lsh (MinHashLSH): The LSH object.
+        minhashes (Dict[str, Tuple[MinHash, str, str, str]]): The dictionary of MinHashes.
+        keyword (str): The keyword to search for.
+        signature_size (int, optional): The size of the MinHash signature.
+        n_gram (int, optional): The n-gram size for the MinHash.
+        top_n (int, optional): The number of top results to return.
+
+    Returns:
+        Dict[str, Dict[str, List[str]]]: A dictionary containing the top similar values.
+        Example:{
+        'table_name1': {
+            'column_name1': ['value1', 'value2', 'value3'],
+            'column_name2': ['value4', 'value5']
+        },
+        'table_name2': {
+    """
+    query_minhash = create_minhash(signature_size, keyword, n_gram)
+    results = lsh.query(query_minhash)
+    similarities = [
+        (result, jaccard_similarity(query_minhash, minhashes[result][0]))
+        for result in results
+    ]
+    similarities = sorted(similarities, key=lambda x: x[1], reverse=True)[:top_n]
+
+    similar_values_trimmed: Dict[str, Dict[str, List[str]]] = {}
+    for result, similarity in similarities:
+        table_name, column_name, value = minhashes[result][1:] #type: ignore
+        if table_name not in similar_values_trimmed:
+            similar_values_trimmed[table_name] = {}
+        if column_name not in similar_values_trimmed[table_name]:
+            similar_values_trimmed[table_name][column_name] = []
+        similar_values_trimmed[table_name][column_name].append(value)
+
+    return similar_values_trimmed