feat: tags for expenses, filters by tags in statistics

services/statistics/expenses_statistics.py

@@ -1,7 +1,10 @@
 from datetime import date
-from sqlalchemy import select, func, Subquery, cast
+
+from sqlalchemy import select, func, Subquery, cast, CTE
 from sqlalchemy.dialects.postgresql import TIMESTAMP
-from models import PaymentRecord, Expense
+
+from models import PaymentRecord, Expense, expenses_expense_tags
+from schemas.statistics import CommonProfitFilters
 from services.base import BaseService
 from services.statistics.common import generate_date_range
 
@@ -10,19 +13,9 @@ class ExpensesStatisticsService(BaseService):
     date_from: date
     date_to: date
 
-    def _get_expenses_sub(self, model, date_column, amount_column) -> Subquery:
-        all_dates = generate_date_range(self.date_from, self.date_to, ["expenses"])
-
-        expenses = (
-            select(
-                func.sum(getattr(model, amount_column)).label("expenses"),
-                cast(getattr(model, date_column), TIMESTAMP(timezone=False)).label("date"),
-            )
-            .group_by("date")
-            .subquery()
-        )
-
-        expenses_with_gaps_filled = (
+    @staticmethod
+    def _fill_date_gaps(expenses: Subquery, all_dates: CTE) -> Subquery:
+        return (
             select(
                 all_dates.c.date,
                 (all_dates.c.expenses + func.coalesce(expenses.c.expenses, 0)).label("expenses"),
@@ -31,9 +24,53 @@ class ExpensesStatisticsService(BaseService):
             .order_by(all_dates.c.date)
             .subquery()
         )
-        return expenses_with_gaps_filled
 
-    def _apply_expenses(self, deals_by_dates: Subquery, expenses_subquery: Subquery):
+    def _get_payment_records_sub(self) -> Subquery:
+        all_dates = generate_date_range(self.date_from, self.date_to, ["expenses"])
+
+        expenses = (
+            select(
+                func.sum(PaymentRecord.amount).label("expenses"),
+                cast(PaymentRecord.start_date, TIMESTAMP(timezone=False)).label("date"),
+            )
+            .group_by("date")
+            .subquery()
+        )
+
+        expenses_with_filled_gaps = self._fill_date_gaps(expenses, all_dates)
+        return expenses_with_filled_gaps
+
+    def _get_additional_expenses_sub(self, tag_id: int) -> Subquery:
+        all_dates = generate_date_range(self.date_from, self.date_to, ["expenses"])
+
+        expenses = (
+            select(Expense)
+        )
+
+        if tag_id != -1:
+            expenses = (
+                expenses
+                .join(expenses_expense_tags)
+                .where(expenses_expense_tags.c.expense_tag_id == tag_id)
+            )
+
+        expenses = expenses.subquery()
+
+        expenses = (
+            select(
+                func.sum(expenses.c.amount).label("expenses"),
+                cast(expenses.c.spent_date, TIMESTAMP(timezone=False)).label("date"),
+            )
+            .where(expenses.c.spent_date.between(self.date_from, self.date_to))
+            .group_by("date")
+            .subquery()
+        )
+
+        expenses_with_filled_gaps = self._fill_date_gaps(expenses, all_dates)
+        return expenses_with_filled_gaps
+
+    @staticmethod
+    def _apply_expenses(deals_by_dates: Subquery, expenses_subquery: Subquery):
         return (
             select(
                 deals_by_dates.c.date,
@@ -46,15 +83,15 @@ class ExpensesStatisticsService(BaseService):
             .join(expenses_subquery, expenses_subquery.c.date == deals_by_dates.c.date)
         )
 
-    def apply_expenses(self, date_from: date, date_to: date, deals_by_dates: Subquery):
-        self.date_from, self.date_to = date_from, date_to
+    def apply_expenses(self, filters: CommonProfitFilters, deals_by_dates: Subquery):
+        self.date_from, self.date_to = filters.date_range
 
         # Apply salary expenses
-        salary_expenses = self._get_expenses_sub(PaymentRecord, "start_date", "amount")
+        salary_expenses = self._get_payment_records_sub()
         deals_by_dates = self._apply_expenses(deals_by_dates, salary_expenses)
 
         # Apply additional expenses
-        additional_expenses = self._get_expenses_sub(Expense, "spent_date", "amount")
+        additional_expenses = self._get_additional_expenses_sub(filters.tag_id)
         deals_by_dates = self._apply_expenses(deals_by_dates, additional_expenses)
 
         return deals_by_dates