get_joined and get_multi_joined methods created

igorbenav · igorbenav · commit 531f96be2ff6 · 2023-11-21T03:11:15.000-03:00
diff --git a/src/app/crud/crud_base.py b/src/app/crud/crud_base.py
@@ -2,14 +2,16 @@
 from datetime import datetime
 
 from pydantic import BaseModel
-from sqlalchemy import select, update, delete, func, and_
+from sqlalchemy import select, update, delete, func, and_, inspect
+from sqlalchemy.sql import Join
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.engine.row import Row
 
-from app.core.models import TimestampModel
 from .helper import (
     _extract_matching_columns_from_schema, 
-    _extract_matching_columns_from_kwargs
+    _extract_matching_columns_from_kwargs,
+    _auto_detect_join_condition,
+    _add_column_with_prefix
 )
 
 ModelType = TypeVar("ModelType")
@@ -191,7 +193,224 @@ async def get_multi(
         total_count = await self.count(db=db, **kwargs)
 
         return {"data": data, "total_count": total_count}
-        
+    
+    async def get_joined(
+            self,
+            db: AsyncSession,
+            join_model: Type[ModelType],
+            join_prefix: str | None = None,
+            join_on: Union[Join, None] = None,
+            schema_to_select: Union[Type[BaseModel], List, None] = None,
+            join_schema_to_select: Union[Type[BaseModel], List, None] = None,
+            join_type: str = "left",
+            **kwargs
+    ) -> Dict | None:
+        """
+        Fetches a single record with a join on another model. If 'join_on' is not provided, the method attempts 
+        to automatically detect the join condition using foreign key relationships.
+
+        Parameters
+        ----------
+        db : AsyncSession
+            The SQLAlchemy async session.
+        join_model : Type[ModelType]
+            The model to join with.
+        join_prefix : Optional[str]
+            Optional prefix to be added to all columns of the joined model. If None, no prefix is added.
+        join_on : Join, optional
+            SQLAlchemy Join object for specifying the ON clause of the join. If None, the join condition is 
+            auto-detected based on foreign keys.
+        schema_to_select : Union[Type[BaseModel], List, None], optional
+            Pydantic schema for selecting specific columns from the primary model.
+        join_schema_to_select : Union[Type[BaseModel], List, None], optional
+            Pydantic schema for selecting specific columns from the joined model.
+        join_type : str, default "left"
+            Specifies the type of join operation to perform. Can be "left" for a left outer join or "inner" for an inner join.
+        kwargs : dict
+            Filters to apply to the query.
+
+        Returns
+        -------
+        Dict | None
+            The fetched database row or None if not found.
+
+        Examples
+        --------
+        Simple example: Joining User and Tier models without explicitly providing join_on
+        ```python
+        result = await crud_user.get_joined(
+            db=session,
+            join_model=Tier,
+            schema_to_select=UserSchema,
+            join_schema_to_select=TierSchema
+        )
+        ```
+
+        Complex example: Joining with a custom join condition, additional filter parameters, and a prefix
+        ```python
+        from sqlalchemy import and_
+        result = await crud_user.get_joined(
+            db=session,
+            join_model=Tier,
+            join_prefix="tier_",
+            join_on=and_(User.tier_id == Tier.id, User.is_superuser == True),
+            schema_to_select=UserSchema,
+            join_schema_to_select=TierSchema,
+            username="john_doe"
+        )
+        ```
+
+        Return example: prefix added, no schema_to_select or join_schema_to_select
+        ```python
+        {
+            "id": 1,
+            "name": "John Doe",
+            "username": "john_doe",
+            "email": "johndoe@example.com",
+            "hashed_password": "hashed_password_example",
+            "profile_image_url": "https://profileimageurl.com/default.jpg",
+            "uuid": "123e4567-e89b-12d3-a456-426614174000",
+            "created_at": "2023-01-01T12:00:00",
+            "updated_at": "2023-01-02T12:00:00",
+            "deleted_at": null,
+            "is_deleted": false,
+            "is_superuser": false,
+            "tier_id": 2,
+            "tier_name": "Premium",
+            "tier_created_at": "2022-12-01T10:00:00",
+            "tier_updated_at": "2023-01-01T11:00:00"
+        }
+        ```
+        """
+        if join_on is None:
+            join_on = _auto_detect_join_condition(self._model, join_model)
+
+        primary_select = _extract_matching_columns_from_schema(model=self._model, schema=schema_to_select)
+        join_select = []
+
+        if join_schema_to_select:
+            columns = _extract_matching_columns_from_schema(model=join_model, schema=join_schema_to_select)
+        else:
+            columns = inspect(join_model).c
+            
+        for column in columns:
+            labeled_column = _add_column_with_prefix(column, join_prefix)
+            if f"{join_prefix}{column.name}" not in [col.name for col in primary_select]:
+                join_select.append(labeled_column)
+
+        if join_type == "left":
+            stmt = select(*primary_select, *join_select).outerjoin(join_model, join_on)
+        elif join_type == "inner":
+            stmt = select(*primary_select, *join_select).join(join_model, join_on)
+        else:
+            raise ValueError(f"Invalid join type: {join_type}. Only 'left' or 'inner' are valid.")
+
+        for key, value in kwargs.items():
+            if hasattr(self._model, key):
+                print(self._model)
+                stmt = stmt.where(getattr(self._model, key) == value)
+
+        db_row = await db.execute(stmt)
+        result = db_row.first()
+        if result:
+            result = dict(result._mapping)
+
+        return result
+    
+    async def get_multi_joined(
+            self,
+            db: AsyncSession,
+            join_model: Type[ModelType],
+            join_prefix: str | None = None,
+            join_on: Union[Join, None] = None,
+            schema_to_select: Union[Type[BaseModel], List[Type[BaseModel]], None] = None,
+            join_schema_to_select: Union[Type[BaseModel], List[Type[BaseModel]], None] = None,
+            join_type: str = "left",
+            offset: int = 0,
+            limit: int = 100,
+            **kwargs: Any
+    ) -> Dict[str, Any]:
+        """
+        Fetch multiple records with a join on another model, allowing for pagination.
+
+        Parameters
+        ----------
+        db : AsyncSession
+            The SQLAlchemy async session.
+        join_model : Type[ModelType]
+            The model to join with.
+        join_prefix : Optional[str]
+            Optional prefix to be added to all columns of the joined model. If None, no prefix is added.
+        join_on : Join, optional
+            SQLAlchemy Join object for specifying the ON clause of the join. If None, the join condition is 
+            auto-detected based on foreign keys.
+        schema_to_select : Union[Type[BaseModel], List[Type[BaseModel]], None], optional
+            Pydantic schema for selecting specific columns from the primary model.
+        join_schema_to_select : Union[Type[BaseModel], List[Type[BaseModel]], None], optional
+            Pydantic schema for selecting specific columns from the joined model.
+        join_type : str, default "left"
+            Specifies the type of join operation to perform. Can be "left" for a left outer join or "inner" for an inner join.
+        offset : int, default 0
+            The offset (number of records to skip) for pagination.
+        limit : int, default 100
+            The limit (maximum number of records to return) for pagination.
+        kwargs : dict
+            Filters to apply to the primary query.
+
+        Returns
+        -------
+        Dict[str, Any]
+            A dictionary containing the fetched rows under 'data' key and total count under 'total_count'.
+
+        Examples
+        --------
+        # Fetching multiple User records joined with Tier records, using left join
+        users = await crud_user.get_multi_joined(
+            db=session,
+            join_model=Tier,
+            join_prefix="tier_",
+            schema_to_select=UserSchema,
+            join_schema_to_select=TierSchema,
+            offset=0,
+            limit=10
+        )
+        """
+        if join_on is None:
+            join_on = _auto_detect_join_condition(self._model, join_model)
+
+        primary_select = _extract_matching_columns_from_schema(model=self._model, schema=schema_to_select)
+        join_select = []
+
+        if join_schema_to_select:
+            columns = _extract_matching_columns_from_schema(model=join_model, schema=join_schema_to_select)
+        else:
+            columns = inspect(join_model).c
+
+        for column in columns:
+            labeled_column = _add_column_with_prefix(column, join_prefix)
+            if f"{join_prefix}{column.name}" not in [col.name for col in primary_select]:
+                join_select.append(labeled_column)
+
+        if join_type == "left":
+            stmt = select(*primary_select, *join_select).outerjoin(join_model, join_on)
+        elif join_type == "inner":
+            stmt = select(*primary_select, *join_select).join(join_model, join_on)
+        else:
+            raise ValueError(f"Invalid join type: {join_type}. Only 'left' or 'inner' are valid.")
+
+        for key, value in kwargs.items():
+            if hasattr(self._model, key):
+                stmt = stmt.where(getattr(self._model, key) == value)
+
+        stmt = stmt.offset(offset).limit(limit)
+
+        db_rows = await db.execute(stmt)
+        data = [dict(row._mapping) for row in db_rows]
+
+        total_count = await self.count(db=db, **kwargs)
+
+        return {"data": data, "total_count": total_count}
+
     async def update(
             self, 
             db: AsyncSession, 
diff --git a/src/app/crud/helper.py b/src/app/crud/helper.py
@@ -1,4 +1,9 @@
-from typing import Any, List, Type, Union
+from typing import Any, List, Type, Union, Optional
+from sqlalchemy import inspect
+from sqlalchemy.orm import DeclarativeMeta
+from sqlalchemy.sql import ColumnElement
+from sqlalchemy.sql.schema import Column
+from sqlalchemy.sql.elements import Label
 
 from pydantic import BaseModel
 
@@ -53,3 +58,61 @@ def _extract_matching_columns_from_column_names(model: Type[Base], column_names:
             column_list.append(getattr(model, column_name))
 
     return column_list
+
+def _auto_detect_join_condition(base_model: Type[DeclarativeMeta], join_model: Type[DeclarativeMeta]) -> Optional[ColumnElement]:
+    """
+    Automatically detects the join condition for SQLAlchemy models based on foreign key relationships.
+    This function scans the foreign keys in the base model and tries to match them with columns in the join model.
+
+    Parameters
+    ----------
+    base_model : Type[DeclarativeMeta]
+        The base SQLAlchemy model from which to join.
+    join_model : Type[DeclarativeMeta]
+        The SQLAlchemy model to join with the base model.
+
+    Returns
+    -------
+    Optional[ColumnElement]
+        A SQLAlchemy ColumnElement representing the join condition, if successfully detected.
+
+    Raises
+    ------
+    ValueError
+        If the join condition cannot be automatically determined, a ValueError is raised.
+
+    Example
+    -------
+    # Assuming User has a foreign key reference to Tier:
+    join_condition = auto_detect_join_condition(User, Tier)
+    """
+    fk_columns = [col for col in inspect(base_model).c if col.foreign_keys]
+    join_on = next(
+        (base_model.__table__.c[col.name] == join_model.__table__.c[list(col.foreign_keys)[0].column.name]
+         for col in fk_columns if list(col.foreign_keys)[0].column.table == join_model.__table__),
+        None
+    )
+
+    if join_on is None:
+        raise ValueError("Could not automatically determine join condition. Please provide join_on.")
+
+    return join_on
+
+def _add_column_with_prefix(column: Column, prefix: Optional[str]) -> Label:
+    """
+    Creates a SQLAlchemy column label with an optional prefix.
+
+    Parameters
+    ----------
+    column : Column
+        The SQLAlchemy Column object to be labeled.
+    prefix : Optional[str]
+        An optional prefix to prepend to the column's name.
+
+    Returns
+    -------
+    Label
+        A labeled SQLAlchemy Column object.
+    """
+    column_label = f"{prefix}{column.name}" if prefix else column.name
+    return column.label(column_label)
diff --git a/src/app/models/user.py b/src/app/models/user.py
@@ -2,7 +2,7 @@
 import uuid as uuid_pkg
 from datetime import datetime
 
-from sqlalchemy import String, DateTime
+from sqlalchemy import String, DateTime, ForeignKey
 from sqlalchemy.orm import Mapped, mapped_column
 
 from app.core.database import Base
@@ -31,4 +31,4 @@ class User(Base):
     is_deleted: Mapped[bool] = mapped_column(default=False, index=True)
     is_superuser: Mapped[bool] = mapped_column(default=False)
 
-    tier_id: Mapped[int | None] = mapped_column(index=True, default=None, init=False)
+    tier_id: Mapped[int | None] = mapped_column(ForeignKey('tier.id'), index=True, default=None, init=False)
