Task and backend support added for external blocks

backend/entityservice/database/insertions.py

@@ -47,7 +47,7 @@ def insert_dataprovider(cur, auth_token, project_id):
 
 def insert_blocking_metadata(db, dp_id, blocks):
     """
-    Insert a new entry into the blocks table.
+    Insert new entries into the blocks table.
 
     :param blocks: A dict mapping block id to the number of encodings per block.
     """
@@ -200,6 +200,24 @@ def update_encoding_metadata(db, clks_filename, dp_id, state):
         ])
 
 
+def update_blocks_state(db, dp_id, blocks, state):
+    assert state in {'pending', 'ready', 'error'}
+    sql_query = """
+        UPDATE blocks
+        SET
+          state = %s
+        WHERE
+          dp = %s AND
+          block_name in %s
+        """
+
+    with db.cursor() as cur:
+        cur.execute(sql_query, [
+            state,
+            dp_id,
+            tuple(blocks)
+        ])
+
 def update_encoding_metadata_set_encoding_size(db, dp_id, encoding_size):
     sql_query = """
         UPDATE uploads
@@ -278,6 +296,18 @@ def update_project_mark_all_runs_failed(conn, project_id):
         cur.execute(sql_query, [project_id])
 
 
+def update_dataprovider_uploaded_state(conn, project_id, dp_id, state):
+    with conn.cursor() as cur:
+        sql_query = """
+            UPDATE dataproviders SET
+              uploaded = %s
+            WHERE
+              id = %s AND
+              project = %s
+            """
+        cur.execute(sql_query, [state, dp_id, project_id])
+
+
 def mark_project_deleted(db, project_id):
     with db.cursor() as cur:
         sql_query = """

backend/entityservice/object_store.py

@@ -6,11 +6,12 @@
 logger = get_logger('objectstore')
 
 
-def connect_to_object_store():
+def connect_to_object_store(credentials=None):
     mc = minio.Minio(
         config.MINIO_SERVER,
         config.MINIO_ACCESS_KEY,
         config.MINIO_SECRET_KEY,
+        credentials=credentials,
         secure=False
     )
     logger.debug("Connected to minio")
@@ -45,3 +46,12 @@ def create_bucket(minio_client, bucket):
             minio_client.make_bucket(bucket)
         except minio.error.BucketAlreadyOwnedByYou:
             logger.info("The bucket {} was already created.".format(bucket))
+
+
+def stat_and_stream_object(bucket_name, object_name, credentials=None):
+    mc = connect_to_object_store(credentials)
+    logger.debug("Checking object exists in object store")
+    stat = mc.stat_object(bucket_name=bucket_name, object_name=object_name)
+    logger.debug("Retrieving file from object store")
+    response = mc.get_object(bucket_name=bucket_name, object_name=object_name)
+    return stat, response

backend/entityservice/tasks/__init__.py

@@ -6,7 +6,7 @@
 from entityservice.tasks.pre_run_check import check_for_executable_runs
 from entityservice.tasks.assert_valid_run import assert_valid_run
 from entityservice.tasks.run import prerun_check
-from entityservice.tasks.encoding_uploading import handle_raw_upload, handle_external_data_pull
+from entityservice.tasks.encoding_uploading import handle_raw_upload, pull_external_data_encodings_only, pull_external_data
 from entityservice.tasks.comparing import create_comparison_jobs, compute_filter_similarity, aggregate_comparisons
 from entityservice.tasks.permutation import save_and_permute, permute_mapping_data
 from entityservice.tasks.solver import solver_task

backend/entityservice/tasks/encoding_uploading.py

@@ -1,46 +1,134 @@
-import io
+import json
 
-import opentracing
+from minio.credentials import Credentials, Static
 
 from entityservice.database import *
 from entityservice.encoding_storage import stream_json_clksnblocks, convert_encodings_from_base64_to_binary, \
     store_encodings_in_db, upload_clk_data_binary, include_encoding_id_in_binary_stream
 from entityservice.error_checking import check_dataproviders_encoding, handle_invalid_encoding_data, \
     InvalidEncodingError
-from entityservice.object_store import connect_to_object_store
+from entityservice.object_store import connect_to_object_store, stat_and_stream_object
+from entityservice.serialization import binary_format
 from entityservice.settings import Config
 from entityservice.async_worker import celery, logger
 from entityservice.tasks.base_task import TracedTask
 from entityservice.tasks.pre_run_check import check_for_executable_runs
-from entityservice.utils import fmt_bytes, clks_uploaded_to_project
+from entityservice.utils import fmt_bytes, clks_uploaded_to_project, object_store_upload_path
+
+logger = get_logger()
 
 
 @celery.task(base=TracedTask, ignore_result=True, args_as_tags=('project_id', 'dp_id'))
-def handle_external_data_pull(project_id, dp_id, object_info, credentials, receipt_token, parent_span=None):
+def pull_external_data(project_id, dp_id,
+                                      encoding_object_info, encoding_credentials,
+                                      blocks_object_info, blocks_credentials,
+                                      receipt_token, parent_span=None):
     """
 
+    - pull blocking map into memory, create blocks in db (with dummy counts?)
+    - stream encodings into DB and add encoding + blocks from in memory dict.
+
     """
     log = logger.bind(pid=project_id, dp_id=dp_id)
-    log.info("Pulling data from an object store")
+    with DBConn() as conn:
+        if not check_project_exists(conn, project_id):
+            log.info("Project deleted, stopping immediately")
+            return
 
-    with DBConn() as db:
-        if not check_project_exists(db, project_id):
+        mc = connect_to_object_store(parse_minio_credentials(blocks_credentials))
+
+    log.debug("Pulling blocking information from object store")
+    response = mc.get_object(bucket_name=blocks_object_info['bucket'], object_name=blocks_object_info['path'])
+    encoding_to_block_map = json.load(response)
+
+    log.debug("Counting the blocks")
+    block_sizes = {}
+    for encoding_id in encoding_to_block_map:
+        _blocks = encoding_to_block_map[encoding_id]
+        for block_id in _blocks:
+            block_id = str(block_id)
+            block_sizes[block_id] = block_sizes.setdefault(block_id, 0) + 1
+
+    block_count = len(block_sizes)
+    log.debug(f"Processing {block_count} blocks")
+
+    # stream the encodings
+    bucket_name = encoding_object_info['bucket']
+    object_name = encoding_object_info['path']
+
+    stat, encodings_stream = stat_and_stream_object(bucket_name, object_name, parse_minio_credentials(encoding_credentials))
+    count = int(stat.metadata['X-Amz-Meta-Hash-Count'])
+    size = int(stat.metadata['X-Amz-Meta-Hash-Size'])
+    log.debug(f"Processing {count} encodings")
+    assert count == len(encoding_to_block_map), f"Expected {count} encodings in blocks. Got {len(encoding_to_block_map)}"
+
+    with DBConn() as conn:
+        with opentracing.tracer.start_span('update-metadata-db', child_of=parent_span):
+            update_encoding_metadata_set_encoding_size(conn, dp_id, size)
+            insert_encoding_metadata(conn, None, dp_id, receipt_token, encoding_count=count, block_count=block_count)
+        with opentracing.tracer.start_span('create-block-entries-in-db', child_of=parent_span):
+            log.debug("Adding blocks to db")
+            insert_blocking_metadata(conn, dp_id, block_sizes)
+
+        def encoding_iterator(encoding_stream):
+            binary_formatter = binary_format(size)
+            for encoding_id in range(count):
+                yield (
+                    str(encoding_id),
+                    binary_formatter.pack(encoding_id, encoding_stream.read(size)),
+                    encoding_to_block_map[str(encoding_id)]
+                    )
+
+        with opentracing.tracer.start_span('upload-encodings-to-db', child_of=parent_span):
+            log.debug("Adding encodings and associated blocks to db")
+            try:
+                store_encodings_in_db(conn, dp_id, encoding_iterator(encodings_stream), size)
+            except Exception as e:
+                update_dataprovider_uploaded_state(conn, project_id, dp_id, 'error')
+                log.warning(e)
+
+        with opentracing.tracer.start_span('update-encoding-metadata', child_of=parent_span):
+            update_encoding_metadata(conn, None, dp_id, 'ready')
+            update_blocks_state(conn, dp_id, block_sizes.keys(), 'ready')
+
+
+@celery.task(base=TracedTask, ignore_result=True, args_as_tags=('project_id', 'dp_id'))
+def pull_external_data_encodings_only(project_id, dp_id, object_info, credentials, receipt_token, parent_span=None):
+    """
+
+    """
+    log = logger.bind(pid=project_id, dp_id=dp_id)
+
+    with DBConn() as conn:
+        if not check_project_exists(conn, project_id):
             log.info("Project deleted, stopping immediately")
             return
 
-    mc = connect_to_object_store()
-    log.info("Retrieving metadata from object store")
-    stat = mc.stat_object(bucket_name=object_info['bucket'], object_name=object_info['path'])
-    logger.info(stat.metadata)
+        bucket_name = object_info['bucket']
+        object_name = object_info['path']
+
+
+    log.info("Pulling encoding data from an object store")
+    mc_credentials = parse_minio_credentials(credentials)
+    stat, stream = stat_and_stream_object(bucket_name, object_name, mc_credentials)
+
     count = int(stat.metadata['X-Amz-Meta-Hash-Count'])
     size = int(stat.metadata['X-Amz-Meta-Hash-Size'])
-    log.info("Retrieving file from object store")
-    response = mc.get_object(bucket_name=object_info['bucket'], object_name=object_info['path'])
-    stream = response.stream()
     converted_stream = include_encoding_id_in_binary_stream(stream, size, count)
     upload_clk_data_binary(project_id, dp_id, converted_stream, receipt_token, count, size)
 
 
+def parse_minio_credentials(credentials):
+    if credentials:
+        access_key = credentials['AccessKeyId']
+        secret_key = credentials['SecretAccessKey']
+        session_token = credentials.get('SessionToken', None)
+        mc_credentials = Credentials(provider=Static(access_key, secret_key, session_token))
+    else:
+        mc_credentials = None
+    return mc_credentials
+
+
 @celery.task(base=TracedTask, ignore_result=True, args_as_tags=('project_id', 'dp_id'))
 def handle_raw_upload(project_id, dp_id, receipt_token, parent_span=None):
     """

backend/entityservice/views/project.py

@@ -8,9 +8,9 @@
 
 import entityservice.database as db
 from entityservice.encoding_storage import upload_clk_data_binary, include_encoding_id_in_binary_stream
-from entityservice.tasks import handle_raw_upload, remove_project, handle_external_data_pull
+from entityservice.tasks import handle_raw_upload, remove_project, pull_external_data_encodings_only, pull_external_data
 from entityservice.tracing import serialize_span
-from entityservice.utils import safe_fail_request, get_json, generate_code
+from entityservice.utils import safe_fail_request, get_json, generate_code, object_store_upload_path
 from entityservice.database import DBConn, get_project_column
 from entityservice.views.auth_checks import abort_if_project_doesnt_exist, abort_if_invalid_dataprovider_token, \
     abort_if_invalid_results_token, get_authorization_token_type_or_abort, abort_if_inconsistent_upload
@@ -119,7 +119,7 @@ def project_binaryclks_post(project_id):
 
     log = log.bind(dp_id=dp_id)
     log.info("Receiving CLK data.")
-    receipt_token = None
+    receipt_token = generate_code()
 
     with opentracing.tracer.start_span('upload-clk-data', child_of=parent_span) as span:
         span.set_tag("project_id", project_id)
@@ -328,22 +328,43 @@ def handle_encoding_upload_json(project_id, dp_id, clk_json, receipt_token, uses
     if "encodings" in clk_json and 'file' in clk_json['encodings']:
         # external encodings
         logger.info("External encodings uploaded")
-        object_info = clk_json['encodings']['file']
-        credentials = clk_json['encodings'].get('credentials')
+        encoding_object_info = clk_json['encodings']['file']
+        object_name = encoding_object_info['path']
+        if not object_name.startswith(object_store_upload_path(project_id, dp_id)):
+            safe_fail_request(403, "Provided object store path is not allowed")
+
+        encoding_credentials = clk_json['encodings'].get('credentials')
         # Schedule a background task to pull the encodings from the object store
         # This background task updates the database with encoding metadata assuming
         # that there are no blocks.
         if 'blocks' not in clk_json:
-            handle_external_data_pull.delay(
+            pull_external_data_encodings_only.delay(
                 project_id,
                 dp_id,
-                object_info,
-                credentials,
+                encoding_object_info,
+                encoding_credentials,
                 receipt_token,
                 parent_span=serialize_span(parent_span))
         else:
-            # TODO need to deal with the optional blocks
-            raise NotImplementedError("Don't currently handle combination of external encodings and blocks")
+            # Need to deal with both encodings and blocks
+            if 'file' in clk_json['blocks']:
+                object_name = clk_json['blocks']['file']['path']
+                if not object_name.startswith(object_store_upload_path(project_id, dp_id)):
+                    safe_fail_request(403, "Provided object store path is not allowed")
+                # Blocks are in an external file
+                blocks_object_info = clk_json['blocks']['file']
+                blocks_credentials = clk_json['blocks'].get('credentials')
+                pull_external_data.delay(
+                    project_id,
+                    dp_id,
+                    encoding_object_info,
+                    encoding_credentials,
+                    blocks_object_info,
+                    blocks_credentials,
+                    receipt_token,
+                    parent_span=serialize_span(parent_span))
+            else:
+                raise NotImplementedError("Don't currently handle combination of external encodings and blocks")
 
 
         return