iceaxe 0.7.1__cp313-cp313-macosx_11_0_arm64.whl

iceaxe/__tests__/test_text_search.py

@@ -0,0 +1,287 @@
+from typing import Optional, TypeVar, cast
+
+import pytest
+
+from iceaxe import Field, TableBase, alias, func, select
+from iceaxe.field import DBFieldInfo
+from iceaxe.postgres import LexemePriority, PostgresFullText
+from iceaxe.session import DBConnection
+
+T = TypeVar("T")
+
+
+class Article(TableBase):
+    """Test model for full-text search."""
+
+    id: int = Field(primary_key=True)
+    title: str = Field(postgres_config=PostgresFullText(language="english", weight="A"))
+    content: str = Field(
+        postgres_config=PostgresFullText(language="english", weight="B")
+    )
+    summary: Optional[str] = Field(
+        default=None, postgres_config=PostgresFullText(language="english", weight="C")
+    )
+
+
+@pytest.mark.asyncio
+async def test_basic_text_search(indexed_db_connection: DBConnection):
+    """Test basic text search functionality using query builder."""
+    # Create test data
+    articles = [
+        Article(
+            id=1, title="Python Programming", content="Learn Python programming basics"
+        ),
+        Article(
+            id=2, title="Database Design", content="Python and database design patterns"
+        ),
+        Article(id=3, title="Web Development", content="Building web apps with Python"),
+    ]
+
+    await indexed_db_connection.insert(articles)
+
+    # Search in title only
+    title_vector = func.to_tsvector("english", Article.title)
+    query = func.to_tsquery("english", "python")
+
+    results = await indexed_db_connection.exec(
+        select(Article).where(title_vector.matches(query))
+    )
+    assert len(results) == 1
+    assert results[0].id == 1
+
+    # Search in content only
+    content_vector = func.to_tsvector("english", Article.content)
+    results = await indexed_db_connection.exec(
+        select(Article).where(content_vector.matches(query))
+    )
+    assert len(results) == 3  # All articles mention Python in content
+
+
+@pytest.mark.asyncio
+async def test_complex_text_search(indexed_db_connection: DBConnection):
+    """Test complex text search queries with boolean operators."""
+    articles = [
+        Article(id=1, title="Python Programming", content="Learn programming basics"),
+        Article(id=2, title="Python Advanced", content="Advanced programming concepts"),
+        Article(
+            id=3, title="JavaScript Basics", content="Learn programming with JavaScript"
+        ),
+    ]
+
+    await indexed_db_connection.insert(articles)
+
+    # Test AND operator
+    vector = func.to_tsvector("english", Article.title)
+    query = func.to_tsquery("english", "python & programming")
+    results = await indexed_db_connection.exec(
+        select(Article).where(vector.matches(query))
+    )
+    assert len(results) == 1
+    assert results[0].id == 1
+
+    # Test OR operator
+    query = func.to_tsquery("english", "python | javascript")
+    results = await indexed_db_connection.exec(
+        select(Article).where(vector.matches(query))
+    )
+    assert len(results) == 3
+    assert {r.id for r in results} == {1, 2, 3}
+
+    # Test NOT operator
+    query = func.to_tsquery("english", "programming & !python")
+    results = await indexed_db_connection.exec(
+        select(Article).where(vector.matches(query))
+    )
+    assert len(results) == 0  # No articles have "programming" without "python" in title
+
+
+@pytest.mark.asyncio
+async def test_combined_field_search(indexed_db_connection: DBConnection):
+    """Test searching across multiple fields."""
+    articles = [
+        Article(
+            id=1,
+            title="Python Guide",
+            content="Learn programming basics",
+            summary="A beginner's guide to Python",
+        ),
+        Article(
+            id=2,
+            title="Programming Tips",
+            content="Python best practices",
+            summary="Advanced Python concepts",
+        ),
+    ]
+
+    await indexed_db_connection.insert(articles)
+
+    # Search across all fields using list syntax
+    vector = func.to_tsvector(
+        "english", [Article.title, Article.content, Article.summary]
+    )
+    query = func.to_tsquery("english", "python & guide")
+
+    results = await indexed_db_connection.exec(
+        select(Article).where(vector.matches(query))
+    )
+    assert len(results) == 1
+    assert results[0].id == 1  # Only first article has both "python" and "guide"
+
+    # Test the original concatenation syntax still works
+    vector_concat = (
+        func.to_tsvector("english", Article.title)
+        .concat(func.to_tsvector("english", Article.content))
+        .concat(func.to_tsvector("english", Article.summary))
+    )
+    query = func.to_tsquery("english", "python & guide")
+
+    results_concat = await indexed_db_connection.exec(
+        select(Article).where(vector_concat.matches(query))
+    )
+    assert len(results_concat) == 1
+    assert results_concat[0].id == 1  # Results should be the same with both approaches
+
+
+@pytest.mark.asyncio
+async def test_weighted_text_search(indexed_db_connection: DBConnection):
+    """Test text search with weighted columns."""
+    articles = [
+        Article(
+            id=1,
+            title="Python Guide",  # Weight A
+            content="Basic Python",  # Weight B
+            summary="Python tutorial",  # Weight C
+        ),
+        Article(
+            id=2,
+            title="Programming",
+            content="Python Guide",
+            summary="Guide to programming",
+        ),
+    ]
+
+    await indexed_db_connection.insert(articles)
+
+    # Search with weights
+    vector = (
+        func.setweight(func.to_tsvector("english", Article.title), "A")
+        .concat(func.setweight(func.to_tsvector("english", Article.content), "B"))
+        .concat(func.setweight(func.to_tsvector("english", Article.summary), "C"))
+    )
+    query = func.to_tsquery("english", "python & guide")
+
+    results = await indexed_db_connection.exec(
+        select((Article, alias("ts_rank", func.ts_rank(vector, query))))
+        .where(vector.matches(query))
+        .order_by("ts_rank", direction="DESC"),
+    )
+    assert len(results) == 2
+    # First article should rank higher because "Python Guide" is in title (weight A)
+    assert results[0][0].id == 1
+    assert results[1][0].id == 2
+    assert results[0][1] > results[1][1]  # Check that rank is higher
+
+
+@pytest.mark.asyncio
+async def test_weight_priority_variants(indexed_db_connection: DBConnection):
+    """Test text search using both string literals and LexemePriority enum for weights."""
+    articles = [
+        Article(
+            id=1,
+            title="Python Guide",  # Weight A (string literal)
+            content="Basic Python",  # Weight B (enum)
+            summary="Python tutorial",  # Weight C (enum)
+        ),
+        Article(
+            id=2,
+            title="Programming",
+            content="Python Guide",
+            summary="Guide to programming",
+        ),
+    ]
+
+    # Create a variant of Article using the enum
+    class ArticleWithEnum(TableBase):
+        id: int = Field(primary_key=True)
+        title: str = Field(
+            postgres_config=PostgresFullText(
+                language="english", weight=LexemePriority.HIGHEST
+            )
+        )
+        content: str = Field(
+            postgres_config=PostgresFullText(
+                language="english", weight=LexemePriority.HIGH
+            )
+        )
+        summary: Optional[str] = Field(
+            default=None,
+            postgres_config=PostgresFullText(
+                language="english", weight=LexemePriority.LOW
+            ),
+        )
+
+    # Verify both models can be created and weights are equivalent
+    assert (
+        cast(
+            PostgresFullText,
+            cast(DBFieldInfo, Article.model_fields["title"]).postgres_config,
+        ).weight
+        == cast(
+            PostgresFullText,
+            cast(DBFieldInfo, ArticleWithEnum.model_fields["title"]).postgres_config,
+        ).weight
+        == "A"
+    )
+    assert (
+        cast(
+            PostgresFullText,
+            cast(DBFieldInfo, Article.model_fields["content"]).postgres_config,
+        ).weight
+        == cast(
+            PostgresFullText,
+            cast(DBFieldInfo, ArticleWithEnum.model_fields["content"]).postgres_config,
+        ).weight
+        == "B"
+    )
+    assert (
+        cast(
+            PostgresFullText,
+            cast(DBFieldInfo, Article.model_fields["summary"]).postgres_config,
+        ).weight
+        == cast(
+            PostgresFullText,
+            cast(DBFieldInfo, ArticleWithEnum.model_fields["summary"]).postgres_config,
+        ).weight
+        == "C"
+    )
+
+    await indexed_db_connection.insert(articles)
+
+    vector = (
+        func.setweight(
+            func.to_tsvector("english", Article.title), LexemePriority.HIGHEST
+        )
+        .concat(
+            func.setweight(
+                func.to_tsvector("english", Article.content), LexemePriority.HIGH
+            )
+        )
+        .concat(
+            func.setweight(
+                func.to_tsvector("english", Article.summary), LexemePriority.LOW
+            )
+        )
+    )
+    query = func.to_tsquery("english", "python & guide")
+
+    results = await indexed_db_connection.exec(
+        select((Article, alias("ts_rank", func.ts_rank(vector, query))))
+        .where(vector.matches(query))
+        .order_by("ts_rank", direction="DESC"),
+    )
+
+    assert len(results) == 2
+    # First article should rank higher because "Python Guide" is in title (HIGHEST weight)
+    assert results[0][0].id == 1
+    assert results[1][0].id == 2
+    assert results[0][1] > results[1][1]

iceaxe/alias_values.py

@@ -0,0 +1,67 @@
+from dataclasses import dataclass
+from typing import Generic, TypeVar, cast
+
+T = TypeVar("T")
+
+
+@dataclass(frozen=True, slots=True)
+class Alias(Generic[T]):
+    name: str
+    value: T
+
+    def __str__(self):
+        return self.name
+
+
+def alias(name: str, type: T) -> T:
+    """
+    Creates an alias for a field in raw SQL queries, allowing for type-safe mapping of raw SQL results.
+    This is particularly useful in two main scenarios:
+
+    1. When using raw SQL queries with aliased columns:
+    ```python
+    # Map a COUNT(*) result to an integer
+    query = select(alias("user_count", int)).text(
+        "SELECT COUNT(*) AS user_count FROM users"
+    )
+
+    # Map multiple aliased columns with different types
+    query = select((
+        alias("full_name", str),
+        alias("order_count", int),
+        alias("total_spent", float)
+    )).text(
+        '''
+        SELECT
+            concat(first_name, ' ', last_name) AS full_name,
+            COUNT(orders.id) AS order_count,
+            SUM(orders.amount) AS total_spent
+        FROM users
+        LEFT JOIN orders ON users.id = orders.user_id
+        GROUP BY users.id
+        '''
+    )
+    ```
+
+    2. When combining ORM models with function results:
+    ```python
+    # Select a model alongside a function result
+    query = select((
+        User,
+        alias("name_length", func.length(User.name)),
+        alias("upper_name", func.upper(User.name))
+    ))
+
+    # Use with aggregation functions
+    query = select((
+        User,
+        alias("total_orders", func.count(Order.id))
+    )).join(Order, User.id == Order.user_id).group_by(User.id)
+    ```
+
+    :param name: The name of the alias as it appears in the SQL query's AS clause
+    :param type: Either a Python type to cast the result to (e.g., int, str, float) or
+                a function metadata object (e.g., from func.length())
+    :return: A type-safe alias that can be used in select() statements
+    """
+    return cast(T, Alias(name, type))

iceaxe/base.py

@@ -0,0 +1,350 @@
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    ClassVar,
+    Self,
+    Type,
+    dataclass_transform,
+)
+
+from pydantic import BaseModel, Field as PydanticField
+from pydantic.main import _model_construction
+from pydantic_core import PydanticUndefined
+
+from iceaxe.field import DBFieldClassDefinition, DBFieldInfo, Field
+
+
+@dataclass_transform(kw_only_default=True, field_specifiers=(PydanticField,))
+class DBModelMetaclass(_model_construction.ModelMetaclass):
+    """
+    Metaclass for database model classes that provides automatic field tracking and SQL query generation.
+    Extends Pydantic's model metaclass to add database-specific functionality.
+
+    This metaclass provides:
+    - Automatic field to SQL column mapping
+    - Dynamic field access that returns query-compatible field definitions
+    - Registry tracking of all database model classes
+    - Support for generic model instantiation
+
+    ```python {{sticky: True}}
+    class User(TableBase):  # Uses DBModelMetaclass
+        id: int = Field(primary_key=True)
+        name: str
+        email: str | None
+
+    # Fields can be accessed for queries
+    User.id      # Returns DBFieldClassDefinition
+    User.name    # Returns DBFieldClassDefinition
+
+    # Metaclass handles model registration
+    registered_models = DBModelMetaclass.get_registry()
+    ```
+    """
+
+    _registry: list[Type["TableBase"]] = []
+    _cached_args: dict[Type["TableBase"], dict[str, Any]] = {}
+    is_constructing: bool = False
+
+    def __new__(mcs, name, bases, namespace, **kwargs):
+        """
+        Create a new database model class with proper field tracking.
+        Handles registration of the model and processes any table-specific arguments.
+        """
+        raw_kwargs = {**kwargs}
+
+        mcs.is_constructing = True
+        autodetect = mcs._extract_kwarg(kwargs, "autodetect", True)
+        cls = super().__new__(mcs, name, bases, namespace, **kwargs)
+        mcs.is_constructing = False
+
+        # Allow future calls to subclasses / generic instantiations to reference the same
+        # kwargs as the base class
+        mcs._cached_args[cls] = raw_kwargs
+
+        # If we have already set the class's fields, we should wrap them
+        if hasattr(cls, "__pydantic_fields__"):
+            cls.__pydantic_fields__ = {
+                field: info
+                if isinstance(info, DBFieldInfo)
+                else DBFieldInfo.extend_field(
+                    info,
+                    primary_key=False,
+                    postgres_config=None,
+                    foreign_key=None,
+                    unique=False,
+                    index=False,
+                    check_expression=None,
+                    is_json=False,
+                )
+                for field, info in cls.model_fields.items()
+            }
+
+        # Avoid registering HandlerBase itself
+        if cls.__name__ not in {"TableBase", "BaseModel"} and autodetect:
+            DBModelMetaclass._registry.append(cls)
+
+        return cls
+
+    def __getattr__(self, key: str) -> Any:
+        """
+        Provides dynamic access to model fields as query-compatible definitions.
+        When accessing an undefined attribute, checks if it's a model field and returns
+        a DBFieldClassDefinition if it is.
+
+        :param key: The attribute name to access
+        :return: Field definition or raises AttributeError
+        :raises AttributeError: If the attribute doesn't exist and isn't a model field
+        """
+        if self.is_constructing:
+            return super().__getattr__(key)  # type: ignore
+
+        try:
+            return super().__getattr__(key)  # type: ignore
+        except AttributeError:
+            # Determine if this field is defined within the spec
+            # If so, return it
+            if key in self.model_fields:
+                return DBFieldClassDefinition(
+                    root_model=self,  # type: ignore
+                    key=key,
+                    field_definition=self.model_fields[key],
+                )
+            raise
+
+    @classmethod
+    def get_registry(cls) -> list[Type["TableBase"]]:
+        """
+        Get the set of all registered database model classes.
+
+        :return: Set of registered TableBase classes
+        """
+        return cls._registry
+
+    @classmethod
+    def _extract_kwarg(
+        cls, kwargs: dict[str, Any], key: str, default: Any = None
+    ) -> Any:
+        """
+        Extract a keyword argument from either standard kwargs or pydantic generic metadata.
+        Handles both normal instantiation and pydantic's generic model instantiation.
+
+        :param kwargs: Dictionary of keyword arguments
+        :param key: Key to extract
+        :param default: Default value if key not found
+        :return: Extracted value or default
+        """
+        if key in kwargs:
+            return kwargs.pop(key)
+
+        if "__pydantic_generic_metadata__" in kwargs:
+            origin_model = kwargs["__pydantic_generic_metadata__"]["origin"]
+            if origin_model in cls._cached_args:
+                return cls._cached_args[origin_model].get(key, default)
+
+        return default
+
+    @property
+    def model_fields(self) -> dict[str, DBFieldInfo]:  # type: ignore
+        """
+        Get the dictionary of model fields and their definitions.
+        Overrides the ClassVar typehint from TableBase for proper typing.
+
+        :return: Dictionary of field names to field definitions
+        """
+        return getattr(self, "__pydantic_fields__", {})  # type: ignore
+
+
+class UniqueConstraint(BaseModel):
+    """
+    Represents a UNIQUE constraint in a database table.
+    Ensures that the specified combination of columns contains unique values across all rows.
+
+    ```python {{sticky: True}}
+    class User(TableBase):
+        email: str
+        tenant_id: int
+
+        table_args = [
+            UniqueConstraint(columns=["email", "tenant_id"])
+        ]
+    ```
+    """
+
+    columns: list[str]
+    """
+    List of column names that should have unique values
+    """
+
+
+class IndexConstraint(BaseModel):
+    """
+    Represents an INDEX on one or more columns in a database table.
+    Improves query performance for the specified columns.
+
+    ```python {{sticky: True}}
+    class User(TableBase):
+        email: str
+        last_login: datetime
+
+        table_args = [
+            IndexConstraint(columns=["last_login"])
+        ]
+    ```
+    """
+
+    columns: list[str]
+    """
+    List of column names to create an index on
+    """
+
+
+INTERNAL_TABLE_FIELDS = ["modified_attrs", "modified_attrs_callbacks"]
+
+
+class TableBase(BaseModel, metaclass=DBModelMetaclass):
+    """
+    Base class for all database table models.
+    Provides the foundation for defining database tables using Python classes with
+    type hints and field definitions.
+
+    Features:
+    - Automatic table name generation from class name
+    - Support for custom table names
+    - Tracking of modified fields for efficient updates
+    - Support for unique constraints and indexes
+    - Integration with Pydantic for validation
+
+    ```python {{sticky: True}}
+    class User(TableBase):
+        # Custom table name (optional)
+        table_name = "users"
+
+        # Fields with types and constraints
+        id: int = Field(primary_key=True)
+        email: str = Field(unique=True)
+        name: str
+        is_active: bool = Field(default=True)
+
+        # Table-level constraints
+        table_args = [
+            UniqueConstraint(columns=["email"]),
+            IndexConstraint(columns=["name"])
+        ]
+
+    # Usage in queries
+    query = select(User).where(User.is_active == True)
+    users = await conn.execute(query)
+    ```
+    """
+
+    if TYPE_CHECKING:
+        model_fields: ClassVar[dict[str, DBFieldInfo]]  # type: ignore
+
+    table_name: ClassVar[str] = PydanticUndefined  # type: ignore
+    """
+    Optional custom name for the table
+    """
+
+    table_args: ClassVar[list[UniqueConstraint | IndexConstraint]] = PydanticUndefined  # type: ignore
+    """
+    Table constraints and indexes
+    """
+
+    # Private methods
+    modified_attrs: dict[str, Any] = Field(default_factory=dict, exclude=True)
+    """
+    Dictionary of modified field values since instantiation or the last clear_modified_attributes() call.
+    Used to construct differential update queries.
+    """
+
+    modified_attrs_callbacks: list[Callable[[Self], None]] = Field(
+        default_factory=list, exclude=True
+    )
+    """
+    List of callbacks to be called when the model is modified.
+    """
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        """
+        Track modified attributes when fields are updated.
+        This allows for efficient database updates by only updating changed fields.
+
+        :param name: Attribute name
+        :param value: New value
+        """
+        if name in self.__class__.model_fields:
+            self.modified_attrs[name] = value
+            for callback in self.modified_attrs_callbacks:
+                callback(self)
+        super().__setattr__(name, value)
+
+    def get_modified_attributes(self) -> dict[str, Any]:
+        """
+        Get the dictionary of attributes that have been modified since instantiation
+        or the last clear_modified_attributes() call.
+
+        :return: Dictionary of modified attribute names and their values
+        """
+        return self.modified_attrs
+
+    def clear_modified_attributes(self) -> None:
+        """
+        Clear the tracking of modified attributes.
+        Typically called after successfully saving changes to the database.
+        """
+        self.modified_attrs.clear()
+
+    @classmethod
+    def get_table_name(cls) -> str:
+        """
+        Get the table name for this model.
+        Uses the custom table_name if set, otherwise converts the class name to lowercase.
+
+        :return: Table name to use in SQL queries
+        """
+        if cls.table_name == PydanticUndefined:
+            return cls.__name__.lower()
+        return cls.table_name
+
+    @classmethod
+    def get_client_fields(cls) -> dict[str, DBFieldInfo]:
+        """
+        Get all fields that should be exposed to clients.
+        Excludes internal fields used for model functionality.
+
+        :return: Dictionary of field names to field definitions
+        """
+        return {
+            field: info
+            for field, info in cls.model_fields.items()
+            if field not in INTERNAL_TABLE_FIELDS
+        }
+
+    def register_modified_callback(self, callback: Callable[[Self], None]) -> None:
+        """
+        Register a callback to be called when the model is modified.
+        """
+        self.modified_attrs_callbacks.append(callback)
+
+    def __eq__(self, other: Any) -> bool:
+        """
+        Compare two model instances, ignoring modified_attrs_callbacks.
+        This ensures that two models with the same data but different callbacks are considered equal.
+        """
+        if not isinstance(other, self.__class__):
+            return False
+
+        # Get all fields except modified_attrs_callbacks
+        fields = {
+            field: value
+            for field, value in self.__dict__.items()
+            if field not in INTERNAL_TABLE_FIELDS
+        }
+        other_fields = {
+            field: value
+            for field, value in other.__dict__.items()
+            if field not in INTERNAL_TABLE_FIELDS
+        }
+
+        return fields == other_fields