Merge branch 'master' into ce-use-compound-datasets

Author: alimourey

.github/actions/prepare-breadbox/action.yaml

@@ -12,7 +12,7 @@ runs:
     - name: "Install and configure Poetry"
       uses: snok/install-poetry@v1
       with:
-        version: 2.2.1
+        version: 1.8.2
         virtualenvs-create: true
         virtualenvs-in-project: true
 

breadbox-client/bump_version_and_publish.py

@@ -1,4 +1,5 @@
 #!/usr/bin/env python3
+# force build
 """
 Breadbox Version Bumping and Publishing Script
 =============================================
@@ -181,7 +182,8 @@ def publish():
 
         # Build and publish the package
         print("  Building and publishing package...")
-        subprocess.run(["poetry", "publish", "--build", "--repository", "public-python"], check=True)
+        subprocess.run(["poetry", "build"], check=True)
+        subprocess.run(["poetry", "publish", "--repository", "public-python"], check=True)
         print("  Package published successfully!")
     except Exception as e:
         print(f"Error publishing package: {str(e)}")

breadbox-client/pyproject.toml

@@ -1,16 +1,24 @@
 [tool.poetry]
 name = "breadbox-client"
-version = "4.3.11"
+version = "4.3.19"
 description = "A client library for accessing Breadbox"
 
 authors = []
 
 readme = "README.md"
 packages = [
-    {include = "breadbox_client"},
     {include = "breadbox_facade"},
 ]
-include = ["CHANGELOG.md", "breadbox_client/py.typed", "breadbox_client/**/*"]
+
+include = [
+    { path = "CHANGELOG.md", format = ["sdist", "wheel"] },
+    # because the breadbox_client directory is generated, it's listed in .gitignore
+    # however, this causes poetry to also ignore these files. There doesn't seem to
+    # be a way to override that behavior so explictly tell it to add all the files
+    # to the wheel and sdist files
+    { path = "breadbox_client/py.typed", format = ["sdist", "wheel"] },
+    { path = "breadbox_client/**/*.py", format = ["sdist", "wheel"] }
+]
 
 [tool.poetry.dependencies]
 python = "^3.9"

breadbox/breadbox/compute/celery.py

@@ -1,11 +1,13 @@
 import psutil
 from celery import Celery, Task, signals
 import os
+import uuid
 from logging import getLogger
 from fastapi import HTTPException
 
 from breadbox.logging import GCPExceptionReporter
 from breadbox.celery_task.utils import check_celery
+from breadbox.utils.debug_event_log import log_event, _get_log_filename
 
 from ..config import Settings, get_settings
 from pydantic import ValidationError
@@ -57,32 +59,6 @@ def _get_rss():
     return process.memory_info().rss
 
 
-# from typing import Optional, Any, Dict
-#
-# @signals.task_prerun.connect
-# def task_prerun_handler(
-#    sender:Optional[Any] =None, task_id : Optional[str] =None, task : Optional[Task]=None, args:Optional[Any]=None, kwargs:Optional[Dict[str, Any]]=None, **extras
-# ):
-#    print(
-#        f"[BEFORE] Running task {sender.name} rss:{_get_rss()} ({task_id}) with args={args}, kwargs={kwargs}"
-#    )
-#
-#
-# @signals.task_postrun.connect
-# def task_postrun_handler(
-#    sender:Optional[Any]=None,
-#    task_id : Optional[str]=None,
-#    task: Optional[Task]=None,
-#    args:Optional[Any] = None,
-#    kwargs:Optional[Dict[str, Any]] =None,
-#    retval:Any = None,
-#    state : Any=None,
-#    **extras,
-# ):
-#    print(
-#        f"[AFTER] Finished task {sender.name} rss:{_get_rss()} ({task_id}) with result={retval}, state={state}"
-#    )
-
 try:
     settings = get_settings()
 except ValidationError:
@@ -107,5 +83,33 @@ def _get_rss():
         )
     app.conf.update(**storage_configuration)  # pyright: ignore
 
+# Set up task logging using Celery signals
+@signals.task_prerun.connect
+def task_prerun_handler(task_id, task, *args, **kwargs):
+    log_filename = _get_log_filename()
+    if log_filename:
+        # Generate a readable task name
+        task_name = task.name if hasattr(task, "name") else str(task)
+        # Log task start
+        log_event(log_filename, "start", task_id, {"n": f"Task {task_name}"})
+
+
+@signals.task_success.connect
+def task_success_handler(result, **kwargs):
+    log_filename = _get_log_filename()
+    if log_filename:
+        task_id = kwargs["sender"].request.id
+        # Log task success
+        log_event(log_filename, "end", task_id, {"s": "success"})
+
+
+@signals.task_failure.connect
+def task_failure_handler(task_id, exception, **kwargs):
+    log_filename = _get_log_filename()
+    if log_filename:
+        # Log task failure
+        log_event(log_filename, "end", task_id, {"s": "error", "e": str(exception)})
+
+
 if __name__ == "__main__":
     app.start()

breadbox/breadbox/service/search.py

@@ -3,6 +3,8 @@
 from typing import Dict, List
 
 import pandas as pd
+from sqlalchemy import insert
+
 from ..crud.metadata import cast_tabular_cell_value_type
 from ..crud import dimension_types as types_crud
 
@@ -33,11 +35,12 @@ class PropertyValuePair:
 class MetadataCacheEntry:
     properties_to_index_df: pd.DataFrame
     columns_metadata: Dict[str, ColumnMetadata]
-    label_by_given_id: Dict[str, str]
-    rows_by_index: Dict[str, Dict[str, str]]
+
+    def get_label_for_given_id(self, given_id):
+        return self.properties_to_index_df.loc[given_id, "label"]
 
     def get_properties_dict(self, given_id: str):
-        return self.rows_by_index.get(given_id)
+        return self.properties_to_index_df.loc[given_id].to_dict()
 
 
 class MetadataCache:
@@ -74,19 +77,13 @@ def get(self, dimension_type_name: str):
 
                 columns_metadata = dict(dimension_type.dataset.columns_metadata)
 
-                label_by_given_id = get_dimension_type_metadata_col(
-                    self.db, dimension_type_name=dimension_type.name, col_name="label"
-                )
-
             rows_by_index = {}
             for record in properties_to_index_df.to_records():
                 rows_by_index[record.index] = record
 
             entry = MetadataCacheEntry(
                 properties_to_index_df=properties_to_index_df,
                 columns_metadata=columns_metadata,
-                label_by_given_id=label_by_given_id,
-                rows_by_index=rows_by_index,
             )
             self.cache[dimension_type_name] = entry
 
@@ -140,36 +137,46 @@ def refresh_search_index_for_dimension_type(
         _delete_search_index_records(db, dimension_type)
         log.info("_delete_search_index_records complete")
 
-    dimension_search_index_rows = []
-
     cache_entry = metadata_cache.get(dimension_type.name)
 
-    for given_id in cache_entry.properties_to_index_df.index:
-        for record in get_property_value_pairs_for_given_id(
-            db=db,
-            dimension_type_name=dimension_type.name,
-            given_id=given_id,
-            metadata_cache=metadata_cache,
-        ):
-            # if given_id in cache_entry.dimension_id_by_given_id:
-            dimension_search_index_rows.append(
-                DimensionSearchIndex(
-                    # dimension_id=cache_entry.dimension_id_by_given_id[given_id],
+    def row_generator():
+        for given_id in cache_entry.properties_to_index_df.index:
+            for record in get_property_value_pairs_for_given_id(
+                db=db,
+                dimension_type_name=dimension_type.name,
+                given_id=given_id,
+                metadata_cache=metadata_cache,
+            ):
+                # if given_id in cache_entry.dimension_id_by_given_id:
+                yield dict(
                     property=record.property,
                     value=record.value,
                     group_id=dimension_type.dataset.group_id,
                     dimension_type_name=dimension_type.name,
                     dimension_given_id=given_id,
-                    label=cache_entry.label_by_given_id[given_id],
+                    label=cache_entry.get_label_for_given_id(given_id),
                 )
-            )
+
+    dimension_search_index_row_count = 0
+    for batch in _make_batches(row_generator(), batch_size=1000):
+        db.execute(insert(DimensionSearchIndex), batch)
+        dimension_search_index_row_count += len(batch)
+        f"Wrote batch of {len(batch)} search index records for {dimension_type_name}"
 
     log.info(
-        f"refresh_search_index_for_dimension_type generated {len(dimension_search_index_rows)} search index records for {len(cache_entry.properties_to_index_df.index)} rows in {dimension_type_name}. Writing..."
+        f"Finished writing all {(dimension_search_index_row_count)} search index records for {len(cache_entry.properties_to_index_df.index)} rows in {dimension_type_name}"
     )
 
-    db.bulk_save_objects(dimension_search_index_rows)
-    log.info("refresh_search_index_for_dimension_type complete")
+
+def _make_batches(iterable, batch_size):
+    batch = []
+    for item in iterable:
+        batch.append(item)
+        if len(batch) >= batch_size:
+            yield batch
+            batch = []
+    if len(batch) > 0:
+        yield batch
 
 
 def _get_datatypes_referencing(db, dimension_type_name):

breadbox/breadbox/startup.py

@@ -1,7 +1,12 @@
 import os
 import pathlib
+import time
+import uuid
 
 from fastapi.routing import APIRouter
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.types import ASGIApp, Receive, Scope, Send
+from .utils.debug_event_log import log_event, _get_log_filename
 
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
@@ -61,6 +66,9 @@ def create_app(settings: Settings):
             scheme_override=scheme_override,
             host_override=host_override,
         )
+    
+    # Add request logging middleware
+    app.add_middleware(RequestLoggingMiddleware)
 
     return app
 
@@ -70,6 +78,63 @@ def ensure_directories_exist(settings):
         os.mkdir(settings.filestore_location)
 
 
+class RequestLoggingMiddleware:
+    def __init__(self, app: ASGIApp):
+        self.app = app
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send):
+        if scope["type"] != "http":
+            await self.app(scope, receive, send)
+            return
+
+        log_filename = _get_log_filename()
+        if not log_filename:
+            # If logging is disabled, just pass through
+            await self.app(scope, receive, send)
+            return
+
+        path = scope.get("path", "unknown")
+        method = scope.get("method", "unknown")
+        request_id = str(uuid.uuid4())
+        
+        # Log request start
+        span_name = f"{method} {path}"
+        log_event(log_filename, "start", request_id, {"n": span_name})
+
+        # Capture the original send function to intercept the response
+        original_send = send
+        response_status = None
+        
+        async def wrapped_send(message):
+            nonlocal response_status
+            
+            if message["type"] == "http.response.start":
+                # Capture the status code when headers are sent
+                response_status = message.get("status", 0)
+            
+            # Check if this is the final response body message
+            is_final_message = (
+                message["type"] == "http.response.body" and 
+                not message.get("more_body", False)
+            )
+            
+            # Pass through to the original send
+            await original_send(message)
+            
+            # If this is the final message, log the completion
+            if is_final_message:
+                log_event(log_filename, "end", request_id, {"s": response_status})
+        
+        # Use the wrapped send function and catch any exceptions
+        try:
+            await self.app(scope, receive, wrapped_send)
+        except Exception as e:
+            # Log the exception as an end event
+            log_event(log_filename, "end", request_id, {"s": 500, "error": str(e)})
+            # Re-raise the exception to maintain the original behavior
+            raise
+
+
 class OverrideMiddleWare:
     def __init__(self, app, scheme_override=None, host_override=None):
         self.app = app

breadbox/breadbox/utils/debug_event_log.py

@@ -0,0 +1,23 @@
+import os
+import datetime
+
+import psutil
+import json
+
+
+def _get_log_filename():
+    return os.environ.get("DEBUG_EVENT_LOG")
+
+def log_event(filename, phase, op_id, context=None):
+    process = psutil.Process(os.getpid())
+    rss = process.memory_info().rss
+    with open(filename, "a") as f:
+        record = {
+            "p": phase,
+            "id": op_id,
+            "t": datetime.datetime.now().isoformat(timespec="seconds"),
+            "m": rss,
+        }
+        if context:
+            record.update(context)
+        f.write(json.dumps(record) + "\n")

breadbox/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "breadbox"
-version = "4.3.11"
+version = "4.3.19"
 description = "A peristent service that stores and fetches datasets"
 authors = ["Jessica Cheng <[email protected]>"]
 packages = [{include = "breadbox"}]