diff --git a/TODO.md b/TODO.md index 0170a3a..b2f54b0 100644 --- a/TODO.md +++ b/TODO.md @@ -28,10 +28,6 @@ Items marked with :fire: are high priority. data. - add more tests where 'auto_commit' is set to False to ensure that commit is not called automatically. -- :fire: Extend `with_count()` to support full multi-segment relationship - paths (mixed reverse FK + forward/reverse M2M) with correct distinct - semantics, zero-row preservation via LEFT JOINs, and reusable join planning - across multiple count annotations. - :fire: perhaps add a `JSON` field type to allow storing JSON data in a field, and an `Object` field type to allow storing arbitrary Python objects? Perhaps a `Binary` field type to allow storing arbitrary binary data? (just uses the @@ -51,6 +47,10 @@ Items marked with :fire: are high priority. - Refactor filter condition handling to use one centralized builder path and keep validation/SQL assembly behavior in sync across code paths (issue #136). +- Support `ForeignKey(..., db_column=...)` consistently across ORM runtime + CRUD/query paths (issue #138). Once closed, rewrite temporary custom-column + regression tests (currently using setup workarounds) to use normal ORM + insert/query flows end-to-end. ## Housekeeping diff --git a/docs/api-reference/query-builder.md b/docs/api-reference/query-builder.md index d4d59fc..c8f4bfc 100644 --- a/docs/api-reference/query-builder.md +++ b/docs/api-reference/query-builder.md @@ -324,9 +324,9 @@ def prefetch_related( **Parameters:** -| Parameter | Type | Description | -| --------- | ----- | --------------------------------------------------- | -| `*paths` | `str` | One or more reverse FK or M2M relationship names | +| Parameter | Type | Description | +| --------- | ----- | ------------------------------------------------ | +| `*paths` | `str` | One or more reverse FK or M2M relationship names | **Returns:** `Self` for method chaining. @@ -391,9 +391,9 @@ def group_by( **Parameters:** -| Parameter | Type | Description | -| --------- | ----- | ------------------------ | -| `*fields` | `str` | Fields to group results | +| Parameter | Type | Description | +| --------- | ----- | ----------------------- | +| `*fields` | `str` | Fields to group results | **Returns:** `Self` for method chaining. @@ -415,8 +415,8 @@ def annotate( **Parameters:** -| Parameter | Type | Description | -| -------------- | --------------- | ----------------------------------- | +| Parameter | Type | Description | +| -------------- | --------------- | ------------------------------------ | | `**aggregates` | `AggregateSpec` | Mapping of output alias to aggregate | **Returns:** `Self` for method chaining. @@ -441,9 +441,9 @@ def having( **Parameters:** -| Parameter | Type | Description | -| -------------- | ------------- | --------------------------------- | -| `**conditions` | `FilterValue` | HAVING conditions with operators | +| Parameter | Type | Description | +| -------------- | ------------- | -------------------------------- | +| `**conditions` | `FilterValue` | HAVING conditions with operators | **Returns:** `Self` for method chaining. @@ -471,24 +471,26 @@ def with_count( **Parameters:** -| Parameter | Type | Default | Description | -| ---------- | ------ | --------- | ----------------------------------- | -| `path` | `str` | *required*| Relationship name to count | -| `alias` | `str` | `"count"` | Output alias for the count column | -| `distinct` | `bool` | `False` | Use `COUNT(DISTINCT ...)` | +| Parameter | Type | Default | Description | +| ---------- | ------ | ---------- | --------------------------------- | +| `path` | `str` | *required* | Relationship path to count | +| `alias` | `str` | `"count"` | Output alias for the count column | +| `distinct` | `bool` | `False` | Use `COUNT(DISTINCT ...)` | **Returns:** `Self` for method chaining. **Raises:** - [`InvalidProjectionError`](exceptions.md#invalidprojectionerror) -- - For invalid aliases, nested paths, or unresolved M2M SQL metadata. + For invalid aliases, invalid terminal relationship types, or + unresolved M2M SQL metadata. - [`InvalidRelationshipError`](exceptions.md#invalidrelationshiperror) -- If `path` is not a valid relationship on the model. > [!NOTE] -> `with_count()` currently supports a single relationship segment -> (for example, `"books"`) and not nested paths like `"books__reviews"`. +> `with_count()` supports nested paths (for example, +> `"author__books"` or `"articles__tags"`), but the terminal segment +> must be to-many (reverse FK or many-to-many). > If no `group_by()` exists, SQLiter automatically groups by current > selected model fields. @@ -814,9 +816,9 @@ def update(self, values: dict[str, Any]) -> int: **Parameters:** -| Parameter | Type | Default | Description | -| --------- | ------------------ | ---------- | ------------------------------- | -| `values` | `dict[str, Any]` | *required* | Field names and their new values | +| Parameter | Type | Default | Description | +| --------- | ---------------- | ---------- | -------------------------------- | +| `values` | `dict[str, Any]` | *required* | Field names and their new values | **Returns:** `int` -- The number of records updated. diff --git a/docs/guide/aggregates.md b/docs/guide/aggregates.md index 6a5b26f..9bc18a9 100644 --- a/docs/guide/aggregates.md +++ b/docs/guide/aggregates.md @@ -115,6 +115,15 @@ rows = ( ) ``` +```python +# Multi-segment path: forward FK + reverse FK terminal +rows = ( + db.select(Book) + .with_count("author__books", alias="author_book_count") + .fetch_dicts() +) +``` + ## Important Notes - Projection queries must use `fetch_dicts()`. @@ -122,5 +131,8 @@ rows = ( `count()` are not available in projection mode. - Aggregate aliases must be non-empty, unique, and must not conflict with model field names. -- `with_count()` currently supports a single relationship segment - (for example `"books"`), not nested paths like `"books__reviews"`. +- `with_count()` supports nested relationship paths (for example, + `"author__books"` or `"articles__tags"`), but the terminal segment + must be a to-many relationship (reverse FK or many-to-many). +- Use `distinct=True` when you need unique terminal-row counts across + fan-out joins. diff --git a/docs/tui-demo/index.md b/docs/tui-demo/index.md index 84a2a82..8c7210b 100644 --- a/docs/tui-demo/index.md +++ b/docs/tui-demo/index.md @@ -107,7 +107,8 @@ Define database constraints like unique fields and foreign key relationships. ### [ORM Features](orm.md) -Advanced ORM features including foreign keys, lazy loading, and reverse relationships. +Advanced ORM features including foreign keys, relationship loading, and +`with_count()` projection reporting. ### [Caching](caching.md) diff --git a/docs/tui-demo/orm.md b/docs/tui-demo/orm.md index 8ef2e17..f9e1bb1 100644 --- a/docs/tui-demo/orm.md +++ b/docs/tui-demo/orm.md @@ -443,6 +443,119 @@ db.close() - Keep ad-hoc reporting SQL aligned with ORM naming - Use the same metadata from descriptor and manager access +## with_count() Basics + +Use `with_count()` to return relationship counts in projection mode. + +```python +# --8<-- [start:with-count-basic] +from sqliter import SqliterDB +from sqliter.orm import BaseDBModel, ForeignKey + +class AuthorCountDemo(BaseDBModel): + name: str + +class BookCountDemo(BaseDBModel): + title: str + author: ForeignKey[AuthorCountDemo] = ForeignKey( + AuthorCountDemo, related_name="books" + ) + +db = SqliterDB(memory=True) +db.create_table(AuthorCountDemo) +db.create_table(BookCountDemo) + +alice = db.insert(AuthorCountDemo(name="Alice")) +bob = db.insert(AuthorCountDemo(name="Bob")) +db.insert(AuthorCountDemo(name="No Books")) + +db.insert(BookCountDemo(title="A1", author=alice)) +db.insert(BookCountDemo(title="A2", author=alice)) +db.insert(BookCountDemo(title="B1", author=bob)) + +rows = ( + db.select(AuthorCountDemo) + .with_count("books", alias="book_count") + .order("name") + .fetch_dicts() +) + +print("Authors with book counts:") +for row in rows: + print(f' {row["name"]}: {row["book_count"]}') + +print("\nRows with zero related records are included.") + +db.close() +# --8<-- [end:with-count-basic] +``` + +### What Happens + +- `with_count("books")` adds a count aggregate over the reverse FK relation +- Query runs in projection mode and returns dictionaries via `fetch_dicts()` +- `LEFT JOIN` semantics keep rows that have no related records + +## with_count() Multi-Segment Paths + +`with_count()` also supports nested relationship paths and distinct counts. + +```python +# --8<-- [start:with-count-multi-segment] +from sqliter import SqliterDB +from sqliter.orm import BaseDBModel, ManyToMany + +class TagCountDemo(BaseDBModel): + name: str + +class ArticleCountDemo(BaseDBModel): + title: str + tags: ManyToMany[TagCountDemo] = ManyToMany( + TagCountDemo, related_name="articles" + ) + +db = SqliterDB(memory=True) +db.create_table(TagCountDemo) +db.create_table(ArticleCountDemo) + +python = db.insert(TagCountDemo(name="python")) +sqlite = db.insert(TagCountDemo(name="sqlite")) +db.insert(TagCountDemo(name="unused")) + +guide = db.insert(ArticleCountDemo(title="Guide")) +tips = db.insert(ArticleCountDemo(title="Tips")) +db.insert(ArticleCountDemo(title="No Tags")) + +guide.tags.add(python, sqlite) +tips.tags.add(python) + +rows = ( + db.select(TagCountDemo) + .with_count("articles__tags", alias="tag_links") + .with_count("articles__tags", alias="unique_tags", distinct=True) + .order("name") + .fetch_dicts() +) + +print("Tags counted across articles__tags:") +for row in rows: + print( + f' {row["name"]}: raw={row["tag_links"]}, ' + f'distinct={row["unique_tags"]}' + ) + +print("\nMulti-segment paths support nested relationship counts.") + +db.close() +# --8<-- [end:with-count-multi-segment] +``` + +### What Happens + +- Path `articles__tags` traverses two relationship segments before counting +- Raw count (`distinct=False`) reflects joined row multiplicity +- Distinct count (`distinct=True`) deduplicates terminal related rows + ## Relationship Filter Traversal Filter records by fields on related models using double underscore syntax. diff --git a/sqliter/query/query.py b/sqliter/query/query.py index 7d2ffd9..04e073b 100644 --- a/sqliter/query/query.py +++ b/sqliter/query/query.py @@ -120,6 +120,14 @@ class JoinInfo: is_nullable: bool +@dataclass(frozen=True) +class WithCountJoinNode: + """Cached projection-join metadata for a with_count() path prefix.""" + + alias: str + model_class: type[BaseDBModel] + + class QueryBuilder(Generic[T]): """Builds and executes database queries for a specific model. @@ -175,6 +183,7 @@ def __init__( self._projection_join_clauses: list[str] = [] self._aggregate_sql_expressions: dict[str, str] = {} self._with_count_targets: dict[str, str] = {} + self._with_count_join_nodes: dict[str, WithCountJoinNode] = {} self._projection_alias_counter: int = 1 if self._fields: @@ -490,25 +499,169 @@ def _add_projection_join(self, join_clause: str) -> None: if join_clause not in self._projection_join_clauses: self._projection_join_clauses.append(join_clause) - def _build_reverse_fk_count_target(self, descriptor: Any) -> str: # noqa: ANN401 - """Build reverse-FK LEFT JOIN and return COUNT target expression.""" + def _resolve_with_count_segment( + self, + current_model: type[BaseDBModel], + segment: str, + path: str, + ) -> tuple[Literal["forward_fk", "reverse_fk", "m2m"], Any]: + """Resolve one with_count() segment to a relationship descriptor.""" + from sqliter.orm.fields import ForeignKey # noqa: PLC0415 + from sqliter.orm.m2m import ( # noqa: PLC0415 + ManyToMany, + ReverseManyToMany, + ) + from sqliter.orm.query import ReverseRelationship # noqa: PLC0415 + + fk_descriptors = getattr(current_model, "fk_descriptors", {}) + fk_descriptor = fk_descriptors.get(segment) + if isinstance(fk_descriptor, ForeignKey): + return ("forward_fk", fk_descriptor) + + descriptor = getattr(current_model, segment, None) + if isinstance(descriptor, ReverseRelationship): + return ("reverse_fk", descriptor) + if isinstance(descriptor, (ManyToMany, ReverseManyToMany)): + return ("m2m", descriptor) + + raise InvalidRelationshipError(path, segment, current_model.__name__) + + def _resolve_with_count_path( + self, path: str + ) -> list[tuple[str, Literal["forward_fk", "reverse_fk", "m2m"], Any]]: + """Resolve all with_count() path segments without mutating state.""" + from sqliter.orm.m2m import ( # noqa: PLC0415 + ManyToMany, + ReverseManyToMany, + ) + + segments = path.split("__") + if not segments or any(not segment for segment in segments): + raise InvalidRelationshipError( + path, path, self.model_class.__name__ + ) + + current_model: type[BaseDBModel] = self.model_class + progressive_path: list[str] = [] + resolved: list[ + tuple[str, Literal["forward_fk", "reverse_fk", "m2m"], Any] + ] = [] + + for segment in segments: + progressive_path.append(segment) + current_path = "__".join(progressive_path) + kind, descriptor = self._resolve_with_count_segment( + current_model, segment, path + ) + if kind == "forward_fk": + if isinstance(descriptor.to_model, str): + msg = ( + "Cannot resolve SQL metadata for relationship " + f"'{path}'." + ) + raise InvalidProjectionError(msg) + current_model = cast("type[BaseDBModel]", descriptor.to_model) + elif kind == "reverse_fk": + current_model = descriptor.from_model + else: + if isinstance(descriptor, ManyToMany) and isinstance( + descriptor.to_model, str + ): + msg = ( + "Cannot resolve SQL metadata for many-to-many " + f"relationship '{path}'." + ) + raise InvalidProjectionError(msg) + if isinstance(descriptor, ReverseManyToMany) and isinstance( + descriptor._from_model, # noqa: SLF001 + str, + ): + msg = ( + "Cannot resolve SQL metadata for many-to-many " + f"relationship '{path}'." + ) + raise InvalidProjectionError(msg) + current_model = _get_prefetch_target_model(descriptor) + resolved.append((current_path, kind, descriptor)) + + _, terminal_kind, _ = resolved[-1] + if terminal_kind == "forward_fk": + msg = ( + "with_count() terminal relationship must be to-many " + "(reverse FK or many-to-many)." + ) + raise InvalidProjectionError(msg) + return resolved + + def _build_forward_fk_with_count_join( + self, + parent_alias: str, + segment: str, + descriptor: Any, # noqa: ANN401 + ) -> WithCountJoinNode: + """Build a forward-FK JOIN for with_count() path traversal.""" + join_alias = self._next_projection_alias() + target_model = cast("type[BaseDBModel]", descriptor.to_model) + target_table = target_model.get_table_name() + fk_column = descriptor.fk_info.db_column or f"{segment}_id" + join_type = "LEFT" if descriptor.fk_info.null else "INNER" + join_clause = ( + f'{join_type} JOIN "{target_table}" AS {join_alias} ' + f'ON {parent_alias}."{fk_column}" = {join_alias}."pk"' + ) + self._add_projection_join(join_clause) + return WithCountJoinNode( + alias=join_alias, + model_class=target_model, + ) + + def _build_reverse_fk_with_count_join( + self, + parent_alias: str, + descriptor: Any, # noqa: ANN401 + ) -> WithCountJoinNode: + """Build a reverse-FK LEFT JOIN for with_count() path traversal.""" join_alias = self._next_projection_alias() related_table = descriptor.from_model.get_table_name() + fk_descriptor = getattr( + descriptor.from_model, "fk_descriptors", {} + ).get(descriptor.fk_field) fk_column = f"{descriptor.fk_field}_id" + if fk_descriptor is not None: + db_column = fk_descriptor.fk_info.db_column + if db_column: + fk_column = db_column join_clause = ( f'LEFT JOIN "{related_table}" AS {join_alias} ' - f'ON {join_alias}."{fk_column}" = t0."pk"' + f'ON {join_alias}."{fk_column}" = {parent_alias}."pk"' ) self._add_projection_join(join_clause) - return f'{join_alias}."pk"' + return WithCountJoinNode( + alias=join_alias, + model_class=descriptor.from_model, + ) + + def _build_m2m_with_count_join( + self, + path: str, + parent_alias: str, + descriptor: Any, # noqa: ANN401 + ) -> WithCountJoinNode: + """Build M2M LEFT JOINs for with_count() path traversal.""" + metadata = descriptor.sql_metadata + if metadata is None: + msg = ( + "Cannot resolve SQL metadata for many-to-many relationship " + f"'{path}'." + ) + raise InvalidProjectionError(msg) - def _build_m2m_count_target(self, metadata: Any) -> str: # noqa: ANN401 - """Build M2M LEFT JOINs and return COUNT target expression.""" junction_alias = self._next_projection_alias() target_alias = self._next_projection_alias() join_junction = ( f'LEFT JOIN "{metadata.junction_table}" AS {junction_alias} ' - f'ON {junction_alias}."{metadata.from_column}" = t0."pk"' + f'ON {junction_alias}."{metadata.from_column}" = ' + f'{parent_alias}."pk"' ) join_target = ( f'LEFT JOIN "{metadata.target_table}" AS {target_alias} ' @@ -516,51 +669,45 @@ def _build_m2m_count_target(self, metadata: Any) -> str: # noqa: ANN401 ) self._add_projection_join(join_junction) self._add_projection_join(join_target) - return f'{target_alias}."pk"' + return WithCountJoinNode( + alias=target_alias, + model_class=_get_prefetch_target_model(descriptor), + ) + + def _build_with_count_target_sql(self, path: str) -> str: + """Build/reuse joins for a with_count() path and return COUNT target.""" + resolved_path = self._resolve_with_count_path(path) + + parent_alias = "t0" + for current_path, kind, descriptor in resolved_path: + join_node = self._with_count_join_nodes.get(current_path) + if join_node is None: + segment = current_path.rsplit("__", maxsplit=1)[-1] + if kind == "forward_fk": + join_node = self._build_forward_fk_with_count_join( + parent_alias, segment, descriptor + ) + elif kind == "reverse_fk": + join_node = self._build_reverse_fk_with_count_join( + parent_alias, descriptor + ) + else: + join_node = self._build_m2m_with_count_join( + path, parent_alias, descriptor + ) + self._with_count_join_nodes[current_path] = join_node + parent_alias = join_node.alias + + return f'{parent_alias}."pk"' def with_count( self, path: str, alias: str = "count", *, distinct: bool = False ) -> Self: """Add a relationship count projection using LEFT JOIN semantics.""" - from sqliter.orm.m2m import ( # noqa: PLC0415 - ManyToMany, - ReverseManyToMany, - ) - from sqliter.orm.query import ReverseRelationship # noqa: PLC0415 - - if "__" in path: - msg = ( - "with_count() currently supports only a single relationship " - "segment." - ) - raise InvalidProjectionError(msg) - clean_alias = self._validate_projection_alias(alias) count_target_sql = self._with_count_targets.get(path) if count_target_sql is None: - descriptor = getattr(self.model_class, path, None) - if descriptor is None: - raise InvalidRelationshipError( - path, path, self.model_class.__name__ - ) - - if isinstance(descriptor, ReverseRelationship): - count_target_sql = self._build_reverse_fk_count_target( - descriptor - ) - elif isinstance(descriptor, (ManyToMany, ReverseManyToMany)): - metadata = descriptor.sql_metadata - if metadata is None: - msg = ( - "Cannot resolve SQL metadata for many-to-many " - f"relationship '{path}'." - ) - raise InvalidProjectionError(msg) - count_target_sql = self._build_m2m_count_target(metadata) - else: - raise InvalidRelationshipError( - path, path, self.model_class.__name__ - ) + count_target_sql = self._build_with_count_target_sql(path) self._with_count_targets[path] = count_target_sql self._aggregates[clean_alias] = AggregateSpec( diff --git a/sqliter/tui/demos/orm.py b/sqliter/tui/demos/orm.py index b8612dc..4111767 100644 --- a/sqliter/tui/demos/orm.py +++ b/sqliter/tui/demos/orm.py @@ -302,6 +302,97 @@ class Article(BaseDBModel): return output.getvalue() +def _run_with_count_basic() -> str: + """Demonstrate single-segment with_count() in projection mode.""" + output = io.StringIO() + + class AuthorCountDemo(BaseDBModel): + name: str + + class BookCountDemo(BaseDBModel): + title: str + author: ForeignKey[AuthorCountDemo] = ForeignKey( + AuthorCountDemo, related_name="books" + ) + + db = SqliterDB(memory=True) + db.create_table(AuthorCountDemo) + db.create_table(BookCountDemo) + + alice = db.insert(AuthorCountDemo(name="Alice")) + bob = db.insert(AuthorCountDemo(name="Bob")) + db.insert(AuthorCountDemo(name="No Books")) + + db.insert(BookCountDemo(title="A1", author=alice)) + db.insert(BookCountDemo(title="A2", author=alice)) + db.insert(BookCountDemo(title="B1", author=bob)) + + rows = ( + db.select(AuthorCountDemo) + .with_count("books", alias="book_count") + .order("name") + .fetch_dicts() + ) + + output.write("Authors with book counts:\n") + for row in rows: + output.write(f" {row['name']}: {row['book_count']}\n") + + output.write("\nRows with zero related records are included.\n") + + db.close() + return output.getvalue() + + +def _run_with_count_multi_segment() -> str: + """Demonstrate multi-segment with_count() and distinct semantics.""" + output = io.StringIO() + + class TagCountDemo(BaseDBModel): + name: str + + class ArticleCountDemo(BaseDBModel): + title: str + tags: ManyToMany[TagCountDemo] = ManyToMany( + TagCountDemo, related_name="articles" + ) + + db = SqliterDB(memory=True) + db.create_table(TagCountDemo) + db.create_table(ArticleCountDemo) + + python = db.insert(TagCountDemo(name="python")) + sqlite = db.insert(TagCountDemo(name="sqlite")) + db.insert(TagCountDemo(name="unused")) + + guide = db.insert(ArticleCountDemo(title="Guide")) + tips = db.insert(ArticleCountDemo(title="Tips")) + db.insert(ArticleCountDemo(title="No Tags")) + + guide.tags.add(python, sqlite) + tips.tags.add(python) + + rows = ( + db.select(TagCountDemo) + .with_count("articles__tags", alias="tag_links") + .with_count("articles__tags", alias="unique_tags", distinct=True) + .order("name") + .fetch_dicts() + ) + + output.write("Tags counted across articles__tags:\n") + for row in rows: + output.write( + f" {row['name']}: raw={row['tag_links']}, " + f"distinct={row['unique_tags']}\n" + ) + + output.write("\nMulti-segment paths support nested relationship counts.\n") + + db.close() + return output.getvalue() + + def _run_select_related_basic() -> str: """Demonstrate eager loading with select_related(). @@ -732,6 +823,22 @@ def get_category() -> DemoCategory: code=extract_demo_code(_run_many_to_many_sql_metadata), execute=_run_many_to_many_sql_metadata, ), + Demo( + id="orm_with_count_basic", + title="with_count() Basics", + description="Count related rows with projection results", + category="orm", + code=extract_demo_code(_run_with_count_basic), + execute=_run_with_count_basic, + ), + Demo( + id="orm_with_count_multi_segment", + title="with_count() Multi-Segment Paths", + description="Count nested relationships with distinct support", + category="orm", + code=extract_demo_code(_run_with_count_multi_segment), + execute=_run_with_count_multi_segment, + ), Demo( id="orm_select_related", title="Eager Loading with select_related()", diff --git a/tests/test_aggregates.py b/tests/test_aggregates.py index 5777e72..e7fa12a 100644 --- a/tests/test_aggregates.py +++ b/tests/test_aggregates.py @@ -3,6 +3,7 @@ from __future__ import annotations import sqlite3 +import time from re import escape from typing import Any @@ -42,6 +43,18 @@ class BookAgg(BaseDBModel): ) +class BookCustomColAgg(BaseDBModel): + """Book model with custom FK db_column for reverse-count regression.""" + + title: str + author: ForeignKey[AuthorAgg] = ForeignKey( + AuthorAgg, + on_delete="CASCADE", + related_name="custom_books", + db_column="author_ref", + ) + + class BookFilterAgg(BaseDBModel): """Book model with forward-FK for projection relationship-filter tests.""" @@ -191,6 +204,47 @@ def test_with_count_reverse_fk_includes_zero_rows( assert usage_by_name == {"Alice": 2, "Bob": 1, "No Books": 0} +def test_with_count_reverse_fk_respects_custom_db_column() -> None: + """with_count reverse-FK joins should use custom FK db_column names.""" + db = SqliterDB(":memory:") + db.create_table(AuthorAgg) + db.create_table(BookCustomColAgg) + + alice = db.insert(AuthorAgg(name="Alice")) + bob = db.insert(AuthorAgg(name="Bob")) + db.insert(AuthorAgg(name="No Books")) + + now = int(time.time()) + table = BookCustomColAgg.get_table_name() + insert_sql = ( + f'INSERT INTO "{table}" ' # noqa: S608 + '("created_at", "updated_at", "title", "author_ref") ' + "VALUES (?, ?, ?, ?)" + ) + conn = db.connect() + cursor = conn.cursor() + # Insert via raw SQL instead of db.insert() so data is written directly + # to "author_ref". This keeps the test focused on with_count read/join + # behavior, independent from ORM write-path db_column mapping. + try: + cursor.execute(insert_sql, (now, now, "CA1", alice.pk)) + cursor.execute(insert_sql, (now, now, "CA2", alice.pk)) + cursor.execute(insert_sql, (now, now, "CB1", bob.pk)) + conn.commit() + + rows = ( + db.select(AuthorAgg) + .with_count("custom_books", alias="usage") + .order("name") + .fetch_dicts() + ) + + usage_by_name = {row["name"]: row["usage"] for row in rows} + assert usage_by_name == {"Alice": 2, "Bob": 1, "No Books": 0} + finally: + db.close() + + def test_with_count_reverse_m2m_includes_zero_rows( relation_db: SqliterDB, ) -> None: @@ -221,6 +275,66 @@ def test_with_count_forward_m2m_includes_zero_rows( assert usage_by_article == {"Guide": 2, "No Tags": 0, "Tips": 1} +def test_with_count_multi_segment_forward_fk_to_reverse_fk( + relation_db: SqliterDB, +) -> None: + """with_count supports forward-FK intermediate hops.""" + rows = ( + relation_db.select(BookAgg) + .with_count("author__books", alias="author_book_count") + .order("title") + .fetch_dicts() + ) + + counts_by_title = {row["title"]: row["author_book_count"] for row in rows} + assert counts_by_title == {"A1": 2, "A2": 2, "B1": 1} + + +def test_with_count_multi_segment_m2m_distinct_semantics( + relation_db: SqliterDB, +) -> None: + """Multi-segment M2M counts should support raw and distinct semantics.""" + raw_rows = ( + relation_db.select(TagAgg) + .with_count("articles__tags", alias="usage") + .order("name") + .fetch_dicts() + ) + raw_usage = {row["name"]: row["usage"] for row in raw_rows} + assert raw_usage == {"python": 3, "sqlite": 2, "unused": 0} + + distinct_rows = ( + relation_db.select(TagAgg) + .with_count("articles__tags", alias="usage", distinct=True) + .order("name") + .fetch_dicts() + ) + distinct_usage = {row["name"]: row["usage"] for row in distinct_rows} + assert distinct_usage == {"python": 2, "sqlite": 2, "unused": 0} + + +def test_with_count_multi_segment_reuses_shared_prefix_joins( + relation_db: SqliterDB, +) -> None: + """with_count should reuse shared path-prefix joins across calls.""" + rows = ( + relation_db.select(TagAgg) + .with_count("articles", alias="article_count", distinct=True) + .with_count("articles__tags", alias="tag_links") + .order("name") + .fetch_dicts() + ) + + usage_by_tag = { + row["name"]: (row["article_count"], row["tag_links"]) for row in rows + } + assert usage_by_tag == { + "python": (2, 3), + "sqlite": (1, 2), + "unused": (0, 0), + } + + def test_with_count_supports_having_filter(relation_db: SqliterDB) -> None: """HAVING can filter on aggregate aliases from with_count().""" rows = ( @@ -454,12 +568,18 @@ def test_having_operator_variants_and_type_validation( def test_with_count_validation_errors(relation_db: SqliterDB) -> None: """with_count() should validate path shape and relationship type.""" - with pytest.raises(InvalidProjectionError, match="single relationship"): + with pytest.raises(InvalidRelationshipError): + relation_db.select(AuthorAgg).with_count("") + + with pytest.raises(InvalidProjectionError, match="terminal relationship"): relation_db.select(AuthorAgg).with_count("books__author") with pytest.raises(InvalidRelationshipError): relation_db.select(AuthorAgg).with_count("unknown") + with pytest.raises(InvalidRelationshipError): + relation_db.select(AuthorAgg).with_count("books__missing") + with pytest.raises(InvalidRelationshipError): relation_db.select(AuthorAgg).with_count("name") @@ -472,6 +592,63 @@ def test_with_count_validation_errors(relation_db: SqliterDB) -> None: relation_db.select(UnresolvedOwnerAgg).with_count("pending") +def test_with_count_rejects_unresolved_forward_fk_target( + relation_db: SqliterDB, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """with_count() should reject unresolved forward FK traversal.""" + monkeypatch.setattr(BookAgg.author, "to_model", "MissingAgg") + + with pytest.raises( + InvalidProjectionError, match="Cannot resolve SQL metadata" + ): + relation_db.select(BookAgg).with_count("author__books") + + +def test_with_count_rejects_unresolved_reverse_m2m_target( + relation_db: SqliterDB, monkeypatch: pytest.MonkeyPatch +) -> None: + """with_count() should reject unresolved reverse M2M traversal.""" + from sqliter.orm.m2m import ReverseManyToMany # noqa: PLC0415 + + descriptor = ReverseManyToMany( + from_model=ArticleAgg, + to_model=TagAgg, + junction_table="tagagg_missingagg", + related_name="broken_reverse", + ) + monkeypatch.setattr(descriptor, "_from_model", "MissingAgg") + monkeypatch.setattr( + TagAgg, + "broken_reverse", + descriptor, + raising=False, + ) + + with pytest.raises( + InvalidProjectionError, match="Cannot resolve SQL metadata" + ): + relation_db.select(TagAgg).with_count("broken_reverse") + + +def test_with_count_rejects_m2m_descriptor_missing_sql_metadata( + relation_db: SqliterDB, monkeypatch: pytest.MonkeyPatch +) -> None: + """with_count() should reject M2M descriptors without SQL metadata.""" + descriptor = ManyToMany(TagAgg) + monkeypatch.setattr( + AuthorAgg, + "broken_tags", + descriptor, + raising=False, + ) + + with pytest.raises( + InvalidProjectionError, match="Cannot resolve SQL metadata" + ): + relation_db.select(AuthorAgg).with_count("broken_tags") + + def test_count_distinct_star_is_rejected(sales_db: SqliterDB) -> None: """COUNT(DISTINCT *) should raise a projection error.""" with pytest.raises(InvalidProjectionError, match="DISTINCT with COUNT"):