Merge branch 'fix/url_safety_net' of https://github.com/Zipstack/unstract into fix/url_safety_net

Author: gaya3-zipstack

backend/usage_v2/helper.py

@@ -1,13 +1,24 @@
 import logging
-from datetime import datetime
+from datetime import date, datetime, time, timedelta
 from typing import Any
 
-from django.db.models import QuerySet, Sum
-from rest_framework.exceptions import APIException
+from django.db.models import Count, QuerySet, Sum
+from django.utils import timezone
+from rest_framework.exceptions import APIException, ValidationError
 
 from .constants import UsageKeys
 from .models import Usage
 
+# Try to import subscription plugin at module level
+try:
+    from pluggable_apps.subscription_v2.subscription_helper import (
+        SubscriptionHelper,
+    )
+
+    SUBSCRIPTION_HELPER = SubscriptionHelper
+except ImportError:
+    SUBSCRIPTION_HELPER = None
+
 logger = logging.getLogger(__name__)
 
 
@@ -98,3 +109,130 @@ def format_usage_response(
             "aggregated_data": aggregated_data,
             "date_range": {"start_date": start_date, "end_date": end_date},
         }
+
+    @staticmethod
+    def _get_subscription_dates(
+        organization_id: str,
+    ) -> tuple[datetime | None, datetime | None]:
+        """Attempt to get trial dates from subscription plugin.
+
+        Assumes subscription plugin returns timezone-aware datetimes or dates
+        that can be safely used without normalization (uniform storage).
+
+        Args:
+            organization_id: The organization identifier
+
+        Returns:
+            Tuple of (start_date, end_date) or (None, None) if not available
+        """
+        # Exit early if subscription plugin is not available
+        if SUBSCRIPTION_HELPER is None:
+            logger.debug("Subscription plugin not available, using fallback dates")
+            return None, None
+
+        # Plugin is available, attempt to get subscription data
+        try:
+            org_plans = SUBSCRIPTION_HELPER.get_subscription(organization_id)
+        except AttributeError as e:
+            logger.warning(f"Subscription plugin missing expected methods: {e}")
+            return None, None
+
+        if (
+            not org_plans
+            or not hasattr(org_plans, "start_date")
+            or not hasattr(org_plans, "end_date")
+        ):
+            return None, None
+
+        # If subscription returns dates, convert to datetime with proper times
+        start_date = org_plans.start_date
+        end_date = org_plans.end_date
+
+        # Handle date objects by converting to start/end of day (assuming UTC)
+        if isinstance(start_date, date) and not isinstance(start_date, datetime):
+            start_date = datetime.combine(start_date, time.min).replace(
+                tzinfo=timezone.utc
+            )
+        if isinstance(end_date, date) and not isinstance(end_date, datetime):
+            end_date = datetime.combine(end_date, time.max).replace(tzinfo=timezone.utc)
+
+        logger.info(f"Using subscription dates for org {organization_id}")
+        return start_date, end_date
+
+    @staticmethod
+    def _calculate_usage_metrics(
+        organization: Any, trial_start_date: datetime, trial_end_date: datetime
+    ) -> tuple[dict, int]:
+        """Calculate usage metrics for the trial period.
+
+        Args:
+            organization: The organization object
+            trial_start_date: Start of trial period
+            trial_end_date: End of trial period
+
+        Returns:
+            Tuple of (aggregated_data, documents_processed)
+        """
+        usage_queryset = Usage.objects.filter(
+            organization=organization,
+            created_at__gte=trial_start_date,
+            created_at__lte=trial_end_date,
+        )
+
+        aggregated_data = usage_queryset.aggregate(
+            total_cost=Sum("cost_in_dollars"),
+            total_tokens=Sum("total_tokens"),
+            unique_runs=Count("run_id", distinct=True),
+            api_calls=Count("id"),
+        )
+
+        documents_processed = (
+            usage_queryset.values("workflow_id", "execution_id").distinct().count()
+        )
+
+        return aggregated_data, documents_processed
+
+    @staticmethod
+    def get_trial_statistics(organization) -> dict[str, Any]:
+        """Get comprehensive trial usage statistics for an organization.
+
+        Args:
+            organization: The organization object for which to retrieve trial statistics.
+                         Must have 'organization_id' and 'created_at' attributes.
+
+        Returns:
+            dict: A dictionary containing comprehensive trial usage statistics with keys:
+                - trial_start_date (str): ISO formatted trial start date
+                - trial_end_date (str): ISO formatted trial end date
+                - total_cost (float): Total cost in dollars during trial period
+                - documents_processed (int): Number of unique document processing operations
+                - api_calls (int): Total number of API calls made
+                - etl_runs (int): Number of unique ETL pipeline runs
+        """
+        trial_start_date, trial_end_date = UsageHelper._get_subscription_dates(
+            organization.organization_id
+        )
+        # Use fallback dates if needed (assuming uniform UTC storage)
+        if not trial_start_date:
+            trial_start_date = organization.created_at
+        if not trial_end_date:
+            # For end date, set to end of day (23:59:59.999999)
+            end_of_trial_date = organization.created_at + timedelta(days=14)
+            trial_end_date = end_of_trial_date.replace(
+                hour=23, minute=59, second=59, microsecond=999999
+            )
+        # Validate trial window - guard against inverted date range
+        if trial_end_date < trial_start_date:
+            raise ValidationError("trial_end_date must be on or after trial_start_date")
+        # Calculate usage metrics
+        aggregated_data, documents_processed = UsageHelper._calculate_usage_metrics(
+            organization, trial_start_date, trial_end_date
+        )
+        return {
+            "trial_start_date": trial_start_date.isoformat(),
+            "trial_end_date": trial_end_date.isoformat(),
+            "total_cost": aggregated_data.get("total_cost", 0) or 0,
+            "documents_processed": documents_processed,
+            "api_calls": aggregated_data.get("api_calls", 0) or 0,
+            "etl_runs": aggregated_data.get("unique_runs", 0) or 0,
+        }

backend/usage_v2/urls.py

@@ -5,6 +5,7 @@
 
 get_token_usage = UsageView.as_view({"get": "get_token_usage"})
 aggregate = UsageView.as_view({"get": UsageView.aggregate.__name__})
+trial_statistics = UsageView.as_view({"get": "get_trial_statistics"})
 usage_list = UsageView.as_view({"get": UsageView.list.__name__})
 usage_detail = UsageView.as_view(
     {
@@ -27,6 +28,11 @@
             aggregate,
             name="aggregate",
         ),
+        path(
+            "trial-statistics/",
+            trial_statistics,
+            name="trial-statistics",
+        ),
         path("<str:pk>/", usage_detail, name="usage_detail"),
     ]
 )

backend/usage_v2/views.py

@@ -5,6 +5,7 @@
 from permissions.permission import IsOrganizationMember
 from rest_framework import status, viewsets
 from rest_framework.decorators import action
+from rest_framework.exceptions import ValidationError
 from rest_framework.filters import OrderingFilter
 from rest_framework.permissions import IsAuthenticated
 from rest_framework.response import Response
@@ -93,3 +94,33 @@ def get_token_usage(self, request: HttpRequest) -> Response:
 
         # Return the result
         return Response(status=status.HTTP_200_OK, data=result)
+
+    @action(detail=False, methods=["get"], url_path="trial-statistics")
+    def get_trial_statistics(self, request: HttpRequest) -> Response:
+        """Retrieves comprehensive trial usage statistics for the current organization.
+
+        Returns:
+            Response: A Response object containing trial usage statistics including:
+                     - trial_start_date: ISO formatted trial start date
+                     - trial_end_date: ISO formatted trial end date
+                     - total_cost: Total cost in dollars during trial period
+                     - documents_processed: Number of unique document processing operations
+                     - api_calls: Total number of API calls made
+                     - etl_runs: Number of unique ETL pipeline runs
+        """
+        user_organization = UserContext.get_organization()
+
+        # Validate organization context
+        if not user_organization:
+            logger.warning("No organization context found for user")
+            raise ValidationError("No organization context available")
+
+        # Get trial statistics from helper
+        trial_stats = UsageHelper.get_trial_statistics(user_organization)
+
+        # Log successful retrieval for audit purposes
+        logger.info(
+            f"Trial statistics retrieved for organization {user_organization.organization_id}"
+        )
+
+        return Response(status=status.HTTP_200_OK, data=trial_stats)

backend/workflow_manager/endpoint_v2/constants.py

@@ -25,6 +25,7 @@ class TableColumns:
 
 class DBConnectionClass:
     SNOWFLAKE = "SnowflakeDB"
+    BIGQUERY = "BigQuery"
 
 
 class Snowflake: