extracted client

worker_pool.py

@@ -3,76 +3,68 @@ import threading
 from typing import List, Any, Callable, Tuple
 import time
 
-# Global thread-local storage for bearer token to avoid passing it around
+# Global thread-local storage for FilevineClient to avoid passing it around
 _thread_local = threading.local()
 
-def get_bearer_token():
-    """Get bearer token from thread local storage"""
-    return getattr(_thread_local, 'bearer', None)
+def get_filevine_client():
+    """Get FilevineClient from thread local storage"""
+    return getattr(_thread_local, 'client', None)
 
-def set_bearer_token(token):
-    """Set bearer token in thread local storage"""
-    _thread_local.bearer = token
+def set_filevine_client(client):
+    """Set FilevineClient in thread local storage"""
+    _thread_local.client = client
 
-def worker_init(bearer_token: str):
-    """Initialize worker with bearer token"""
-    set_bearer_token(bearer_token)
+def worker_init(client: 'FilevineClient'):
+    """Initialize worker with FilevineClient"""
+    set_filevine_client(client)
 
-def process_project(index: int, total: int, project_data: dict, bearer_token: str) -> dict:
+def process_project(index: int, total: int, project_data: dict, client: 'FilevineClient') -> dict:
     """
     Process a single project with all its API calls.
     This is the function that will be executed by workers in parallel.
     """
-    # Set the bearer token for this thread
-    set_bearer_token(bearer_token)
+    # Set the FilevineClient for this thread
+    set_filevine_client(client)
     
-    from app import (
-        fetch_client, 
-        fetch_contacts, 
-        fetch_project_detail, 
-        fetch_form, 
-        fetch_collection,
-        fetch_project_tasks,
-        fetch_project_team,
-        convert_to_pacific_time
-    )
+    from app import convert_to_pacific_time
     
     p = project_data
     pid = (p.get("projectId") or {}).get("native")
     print(f"Working on {pid} ({index}/{total})")
-    c = fetch_client(bearer_token, (p.get("clientId") or {}).get("native"))
-    cs = fetch_contacts(bearer_token, pid)
+    client = get_filevine_client()
+    c = client.fetch_client((p.get("clientId") or {}).get("native"))
+    cs = client.fetch_contacts(pid)
 
     if pid is None:
         return {}
 
     try:
-        detail = fetch_project_detail(bearer_token, pid)
+        detail = client.fetch_project_detail(pid)
     except Exception as e:
         print(f"[WARN] detail fetch failed for {pid}: {e}")
         detail = {}
         
     defendant_one = next((c.get('orgContact', {}) for c in cs if "Defendant" in c.get('orgContact', {}).get('personTypes', [])), {})
     
-    new_file_review = fetch_form(bearer_token, pid, "newFileReview") or {}
-    dates_and_deadlines = fetch_form(bearer_token, pid, "datesAndDeadlines") or {}
-    service_info = fetch_collection(bearer_token, pid, "serviceInfo") or []
-    property_info = fetch_form(bearer_token, pid, "propertyInfo")
-    matter_overview = fetch_form(bearer_token, pid, "matterOverview")
-    fees_and_costs = fetch_form(bearer_token, pid, "feesAndCosts") or {}
-    property_contacts = fetch_form(bearer_token, pid, "propertyContacts") or {}
-    lease_info_np = fetch_form(bearer_token, pid, "leaseInfoNP") or {}
+    new_file_review = client.fetch_form(pid, "newFileReview") or {}
+    dates_and_deadlines = client.fetch_form(pid, "datesAndDeadlines") or {}
+    service_info = client.fetch_collection(pid, "serviceInfo") or []
+    property_info = client.fetch_form(pid, "propertyInfo")
+    matter_overview = client.fetch_form(pid, "matterOverview")
+    fees_and_costs = client.fetch_form(pid, "feesAndCosts") or {}
+    property_contacts = client.fetch_form(pid, "propertyContacts") or {}
+    lease_info_np = client.fetch_form(pid, "leaseInfoNP") or {}
     
     completed_tasks = [{"description": x.get("body"),
                         "completed": convert_to_pacific_time(x.get("completedDate"))}
-                       for x in fetch_project_tasks(bearer_token, pid).get("items")
+                       for x in client.fetch_project_tasks(pid).get("items")
                        if x.get("isCompleted")]
     pending_tasks = [{"description": x.get("body"),
                       "completed": convert_to_pacific_time(x.get("completedDate"))}
-                     for x in fetch_project_tasks(bearer_token, pid).get("items")
+                     for x in client.fetch_project_tasks(pid).get("items")
                      if not x.get("isCompleted")]
     
-    team = fetch_project_team(bearer_token, pid)
+    team = client.fetch_project_team(pid)
     assigned_attorney = next((m.get('fullname') 
                              for m in team 
                              if ('Assigned Attorney' in [r.get('name') for r in m.get('teamOrgRoles')])
@@ -202,13 +194,13 @@ def process_project(index: int, total: int, project_data: dict, bearer_token: st
     
     return row
 
-def process_projects_parallel(projects: List[dict], bearer_token: str, max_workers: int = 9) -> List[dict]:
+def process_projects_parallel(projects: List[dict], client: 'FilevineClient', max_workers: int = 9) -> List[dict]:
     """
     Process projects in parallel using a worker pool.
     
     Args:
         projects: List of project data dictionaries
-        bearer_token: Filevine API bearer token
+        client: FilevineClient instance
         max_workers: Number of concurrent workers (default 20)
         
     Returns:
@@ -216,9 +208,9 @@ def process_projects_parallel(projects: List[dict], bearer_token: str, max_worke
     """
     # Create a thread pool with specified number of workers
     total = len(projects)
-    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers, initializer=worker_init, initargs=(bearer_token,)) as executor:
+    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers, initializer=worker_init, initargs=(client,)) as executor:
         # Submit all tasks to the executor
-        future_to_project = {executor.submit(process_project, indx, total, project, bearer_token): project for indx, project in enumerate(projects)}
+        future_to_project = {executor.submit(process_project, indx, total, project, client): project for indx, project in enumerate(projects)}
         
         # Collect results as they complete
         results = []