Zipstack
diff --git a/‎backend/migrating/v2/unstract_migrations.py‎
Lines changed: 62 additions & 1 deletion b/‎backend/migrating/v2/unstract_migrations.py‎
Lines changed: 62 additions & 1 deletion
diff --git a/‎backend/plugins/workflow_manager/workflow_v2/api_hub_usage_utils.py‎
Lines changed: 51 additions & 28 deletions b/‎backend/plugins/workflow_manager/workflow_v2/api_hub_usage_utils.py‎
Lines changed: 51 additions & 28 deletions
@@ -1,5 +1,66 @@
+"""Migration utilities for Unstract.
+
+This module provides the UnstractMigration class which extends MigrationQuery.
+Additional migrations can be provided through pluggable_apps.migrations_ext
+which will be dynamically loaded if available.
+"""
+
+import logging
+
 from migrating.v2.query import MigrationQuery
 
+logger = logging.getLogger(__name__)
+
 
 class UnstractMigration(MigrationQuery):
-    pass
+    """Migration query class that supports extensible migrations.
+
+    Additional migrations can be loaded from pluggable_apps.migrations_ext
+    if that module is available.
+    """
+
+    def get_public_schema_migrations(self) -> list[dict[str, str]]:
+        """Returns a list of dictionaries containing the schema migration details.
+
+        Returns:
+            list: A list of dictionaries containing the schema migration details.
+        """
+        core_migrations = super().get_public_schema_migrations()
+
+        try:
+            from pluggable_apps.migrations_ext.migrations import (
+                get_extended_public_schema_migrations,
+            )
+
+            return core_migrations + get_extended_public_schema_migrations(self.v2_schema)
+        except ImportError:
+            pass
+
+        return core_migrations
+
+    def get_organization_migrations(
+        self, schema: str, organization_id: str
+    ) -> list[dict[str, str]]:
+        """Returns a list of dictionaries containing the organization migration details.
+
+        Args:
+            schema (str): The name of the schema for the organization.
+            organization_id (str): The ID of the organization.
+
+        Returns:
+            list: A list of dictionaries containing the organization migration details.
+        """
+        core_migrations = super().get_organization_migrations(schema, organization_id)
+
+        try:
+            from pluggable_apps.migrations_ext.migrations import (
+                get_extended_organization_migrations,
+            )
+
+            return core_migrations + get_extended_organization_migrations(
+                self.v2_schema, schema, organization_id
+            )
+        except ImportError:
+            pass
+
+        return core_migrations
@@ -1,61 +1,79 @@
 """API Hub usage tracking utilities for workflow execution.
 
 This module provides a clean interface for API Hub usage tracking in workflows.
-The OSS version provides default no-op implementations that can be overridden
-by enterprise plugins during the build process.
+Usage tracking functionality is loaded via get_plugin if available.
 """
 
 import logging
 
+from plugins import get_plugin
+
 logger = logging.getLogger(__name__)
 
 
 class APIHubUsageUtil:
-    """Utility class for handling API Hub usage tracking in workflows.
-
-    This is the OSS version that provides default no-op implementations.
-    Enterprise builds will replace this file with enhanced functionality.
-    """
+    """Utility class for handling API Hub usage tracking in workflows."""
 
     @staticmethod
     def track_api_hub_usage(
         workflow_execution_id: str,
         workflow_file_execution_id: str,
         organization_id: str | None = None,
     ) -> bool:
-        """Track API hub usage if enterprise plugin is available.
-
-        OSS version: This is a no-op implementation that always returns False.
-        Enterprise version: Will track actual usage in verticals.subscription_usage table.
+        """Track API hub usage if the verticals_usage plugin is available.
 
         Args:
             workflow_execution_id: The workflow execution ID
             workflow_file_execution_id: The file execution ID
             organization_id: Optional organization ID
 
         Returns:
-            bool: False in OSS version (usage tracking not available).
-                  Enterprise version returns True if tracking succeeded.
+            bool: True if tracking succeeded, False otherwise.
         """
-        # OSS version - no usage tracking available
-        logger.debug("API hub usage tracking not available in OSS version")
-        return False
+        verticals_usage_plugin = get_plugin("verticals_usage")
+        if not verticals_usage_plugin:
+            return False
+
+        try:
+            headers_cache = verticals_usage_plugin["headers_cache_class"]()
+            usage_tracker = verticals_usage_plugin["service_class"]()
+
+            api_hub_headers = headers_cache.get_headers(workflow_execution_id)
+
+            if api_hub_headers:
+                return usage_tracker.store_usage(
+                    file_execution_id=workflow_file_execution_id,
+                    api_hub_headers=api_hub_headers,
+                    organization_id=organization_id,
+                )
+            return False
+
+        except Exception as e:
+            logger.error(
+                f"Failed to track API hub usage for execution {workflow_execution_id}: {e}"
+            )
+            return False
 
     @staticmethod
     def extract_api_hub_headers(request_headers: dict) -> dict | None:
         """Extract API hub headers from request headers.
 
-        OSS version: Returns None (no API hub support).
-        Enterprise version: Extracts and normalizes API hub headers.
-
         Args:
             request_headers: Request headers dictionary
 
         Returns:
-            None in OSS version. Enterprise version returns normalized headers.
+            Normalized API hub headers or None if not available.
         """
-        # OSS version - no API hub header extraction
-        return None
+        verticals_usage_plugin = get_plugin("verticals_usage")
+        if not verticals_usage_plugin:
+            return None
+
+        try:
+            usage_tracker = verticals_usage_plugin["service_class"]()
+            return usage_tracker.extract_api_hub_headers_from_request(request_headers)
+        except Exception as e:
+            logger.error(f"Error extracting API hub headers: {e}")
+            return None
 
     @staticmethod
     def cache_api_hub_headers(
@@ -65,16 +83,21 @@ def cache_api_hub_headers(
     ) -> bool:
         """Cache API hub headers for later usage tracking.
 
-        OSS version: No-op implementation.
-        Enterprise version: Caches headers in Redis.
-
         Args:
             execution_id: The execution ID to use as cache key
             headers: Headers to cache
             ttl_seconds: Time-to-live in seconds (default 2 hours)
 
         Returns:
-            False in OSS version. Enterprise version returns True if caching succeeded.
+            True if caching succeeded, False otherwise.
         """
-        # OSS version - no caching available
-        return False
+        verticals_usage_plugin = get_plugin("verticals_usage")
+        if not verticals_usage_plugin:
+            return False
+
+        try:
+            headers_cache = verticals_usage_plugin["headers_cache_class"]()
+            return headers_cache.store_headers(execution_id, headers)
+        except Exception as e:
+            logger.error(f"Error caching API hub headers: {e}")
+            return False