Add info argument to call-back

Change in the API - will pass in information about the DID request that can be used by other parts

Author: gordonwatts

README.md

@@ -11,7 +11,7 @@ ServiceX DID finders take a dataset name and turn them into files to be transfor
 Create an async callback method that `yield`s file info dictionaries. For example:
 
 ```python
-    async def my_callback(did_name: str):
+    async def my_callback(did_name: str, info: Dict[str, Any]):
         for i in range(0, 10):
             yield {
                 'file_path': f"root://atlas-experiment.cern.ch/dataset1/file{i}.root",
@@ -23,6 +23,9 @@ Create an async callback method that `yield`s file info dictionaries. For exampl
 
 Yield the results as you find them - ServiceX will actually start processing the files before your DID lookup is finished using this scheme. Fill in `file_size` in bytes and number of events in `file_events` if you know them.
 
+`info` contains a dict of various info about the request that asked for this DID:
+* `request-id` The request id that has this DID asscoiated. For logging.
+
 In your main script, start off the DID finder with a call similar to:
 
 ```python
@@ -31,6 +34,10 @@ start_did_finder('my_finder', my_callback)
 
 This script should be run in your docker container at start up. Once running, if a user submits a query with the DID name `my_finder://dataset1`, then your callback will be called with `did_name` set to `dataset1`.
 
+### Proper Logging
+
+You might not think logging is important, but... k8... blah blah
+
 ## Left to do
 
 - what is adler32

src/servicex_did_finder_lib/communication.py

@@ -14,7 +14,7 @@
 from .servicex_adaptor import ServiceXAdapter
 
 # The type for the callback method to handle DID's, supplied by the user.
-UserDIDHandler = Callable[[str], AsyncGenerator[Dict[str, Any], None]]
+UserDIDHandler = Callable[[str, Dict[str, Any]], AsyncGenerator[Dict[str, Any], None]]
 
 # Given name, build the RabbitMQ queue name by appending this.
 # This is backed into how ServiceX works - do not change unless it
@@ -25,10 +25,11 @@
 __logging.addHandler(logging.NullHandler())
 
 
-async def run_file_fetch_loop(did: str, servicex: ServiceXAdapter, user_callback: UserDIDHandler):
+async def run_file_fetch_loop(did: str, servicex: ServiceXAdapter, info: Dict[str, Any],
+                              user_callback: UserDIDHandler):
     summary = DIDSummary(did)
     start_time = datetime.now()
-    async for file_info in user_callback(did):
+    async for file_info in user_callback(did, info):
 
         # Track the file, inject back into the system
         summary.add_file(file_info)
@@ -78,9 +79,13 @@ def rabbit_mq_callback(user_callback: UserDIDHandler, channel, method, propertie
         servicex = ServiceXAdapter(did_request['service-endpoint'])
         servicex.post_status_update("DID Request received")
 
+        info = {
+            'request-id': request_id,
+        }
+
         # Process the request and resolve the DID
         try:
-            make_sync(run_file_fetch_loop)(did, servicex, user_callback)
+            make_sync(run_file_fetch_loop)(did, servicex, info, user_callback)
 
         except Exception as e:
             traceback.print_exc()

tests/test_communication.py

@@ -1,11 +1,13 @@
 import argparse
+import json
 from typing import Any, AsyncGenerator, Dict
+from unittest.mock import ANY, MagicMock, patch
+
 import pika
 import pytest
-from unittest.mock import ANY, patch, MagicMock
-import json
-
-from servicex_did_finder_lib.communication import default_command_line_args, init_rabbit_mq, start_did_finder
+from servicex_did_finder_lib.communication import (default_command_line_args,
+                                                   init_rabbit_mq,
+                                                   start_did_finder)
 
 
 class RabbitAdaptor:
@@ -96,7 +98,7 @@ def test_one_file_call(rabbitmq, SXAdaptor):
 
     seen_name = None
 
-    async def my_callback(did_name: str):
+    async def my_callback(did_name: str, info: Dict[str, Any]):
         nonlocal seen_name
         seen_name = did_name
         yield {
@@ -122,7 +124,8 @@ async def my_callback(did_name: str):
 def test_failed_file(rabbitmq, SXAdaptor):
     'Test a callback that fails before any files are sent'
 
-    async def my_callback(did_name: str) -> AsyncGenerator[Dict[str, Any], None]:
+    async def my_callback(did_name: str, info: Dict[str, Any]) \
+            -> AsyncGenerator[Dict[str, Any], None]:
         if False:
             yield {
                 'ops': 'no'
@@ -142,7 +145,8 @@ async def my_callback(did_name: str) -> AsyncGenerator[Dict[str, Any], None]:
 def test_no_files_returned(rabbitmq, SXAdaptor):
     'Test a callback that fails before any files are sent'
 
-    async def my_callback(did_name: str) -> AsyncGenerator[Dict[str, Any], None]:
+    async def my_callback(did_name: str, info: Dict[str, Any]) \
+            -> AsyncGenerator[Dict[str, Any], None]:
         if False:
             yield {
                 'ops': 'no'
@@ -163,7 +167,8 @@ def test_rabbitmq_connection_failure(rabbitmq_fail_once, SXAdaptor):
 
     called = False
 
-    async def my_callback(did_name: str) -> AsyncGenerator[Dict[str, Any], None]:
+    async def my_callback(did_name: str, info: Dict[str, Any]) \
+            -> AsyncGenerator[Dict[str, Any], None]:
         nonlocal called
         called = True
         yield {
@@ -183,7 +188,8 @@ async def my_callback(did_name: str) -> AsyncGenerator[Dict[str, Any], None]:
 def test_auto_args_callback(init_rabbit_callback, simple_argument_parser):
     'If there is a missing argument on the command line it should cause a total failure'
 
-    async def my_callback(did_name: str) -> AsyncGenerator[Dict[str, Any], None]:
+    async def my_callback(did_name: str, info: Dict[str, Any]) \
+            -> AsyncGenerator[Dict[str, Any], None]:
         if False:
             yield {
                 'ops': 'no'
@@ -217,7 +223,8 @@ def test_arg_required(init_rabbit_callback):
 
     args = parser.parse_args(['--rabbit-uri', 'not-really-there'])
 
-    async def my_callback(did_name: str) -> AsyncGenerator[Dict[str, Any], None]:
+    async def my_callback(did_name: str, info: Dict[str, Any]) \
+            -> AsyncGenerator[Dict[str, Any], None]:
         if False:
             yield {
                 'ops': 'no'