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
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'])
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)
# # Now work out if all parties have added their data
if clks_uploaded_to_project(project_id):
logger.info("All parties data present. Scheduling any queued runs")
check_for_executable_runs.delay(project_id, serialize_span(parent_span))
def check_for_executable_runs(project_id, parent_span=None):
"""
This is called when a run is posted (if project is ready for runs), and also
after all dataproviders have uploaded CLKs, and the CLKS are ready.
"""
log = logger.bind(pid=project_id)
log.debug("Checking for runs that need to be executed")
if not clks_uploaded_to_project(project_id, check_data_ready=True):
return
with DBConn() as conn:
try:
check_and_set_project_encoding_size(project_id, conn)
except ValueError as e:
log.warning(e.args[0])
# make sure this error can be exposed to user by marking the run/s as failed
update_project_mark_all_runs_failed(conn, project_id)
return
new_runs = get_created_runs_and_queue(conn, project_id)
log.debug("Progressing run stages")
for qr in new_runs:
# Record that the run has reached a new stage
run_id = qr[0]
progress_stage(conn, run_id)
# commit db changes before scheduling following tasks
log.debug("Creating tasks for {} created runs for project {}".format(len(new_runs), project_id))
for qr in new_runs:
run_id = qr[0]
log.info('Queueing run for computation', run_id=run_id)
prerun_check.delay(project_id, run_id, check_for_executable_runs.get_serialized_span())
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
bucket_name = object_info['bucket']
object_name = object_info['path']
log.info("Pulling encoding data from an object store")
env_credentials = parse_minio_credentials({
'AccessKeyId': config.MINIO_ACCESS_KEY,
'SecretAccessKey': config.MINIO_SECRET_KEY
})
stat, stream = stat_and_stream_object(bucket_name, object_name, env_credentials)
count = int(stat.metadata['X-Amz-Meta-Hash-Count'])
size = int(stat.metadata['X-Amz-Meta-Hash-Size'])
if object_name.endswith('.json'):
encodings_stream = ijson.items(io.BytesIO(stream.data), 'clks.item')
converted_stream = include_encoding_id_in_json_stream(encodings_stream, size, count)
else:
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, parent_span=parent_span)
# # Now work out if all parties have added their data
if clks_uploaded_to_project(project_id):
logger.info("All parties data present. Scheduling any queued runs")
check_for_executable_runs.delay(project_id, serialize_span(parent_span))
def upload_clk_data_binary(project_id, dp_id, raw_stream, count, size=128):
"""
Save the user provided raw CLK data.
"""
receipt_token = generate_code()
filename = Config.BIN_FILENAME_FMT.format(receipt_token)
# Set the state to 'pending' in the bloomingdata table
with DBConn() as conn:
db.insert_encoding_metadata(conn, filename, dp_id, receipt_token,
count)
db.update_encoding_metadata_set_encoding_size(conn, dp_id, size)
logger.info(
f"Storing supplied binary clks of individual size {size} in file: {filename}"
)
num_bytes = count * (size + 6)
logger.debug(
"Directly storing binary file with index, base64 encoded CLK, popcount"
)
# Upload to object store
logger.info(
f"Uploading {count} binary encodings to object store. Total size: {fmt_bytes(num_bytes)}"
)
parent_span = g.flask_tracer.get_span()
with opentracing.tracer.start_span('save-to-minio',
child_of=parent_span) as span:
mc = connect_to_object_store()
try:
mc.put_object(Config.MINIO_BUCKET,
filename,
data=raw_stream,
length=num_bytes)
except (minio.error.InvalidSizeError, minio.error.InvalidArgumentError,
minio.error.ResponseError):
logger.info(
"Mismatch between expected stream length and header info")
raise ValueError(
"Mismatch between expected stream length and header info")
with opentracing.tracer.start_span('update-database',
child_of=parent_span) as span:
with DBConn() as conn:
db.update_encoding_metadata(conn, filename, dp_id, 'ready')
db.set_dataprovider_upload_state(conn, dp_id, True)
# Now work out if all parties have added their data
if clks_uploaded_to_project(project_id):
logger.info("All parties data present. Scheduling any queued runs")
check_for_executable_runs.delay(project_id,
serialize_span(parent_span))
return receipt_token
def handle_raw_upload(project_id, dp_id, receipt_token, parent_span=None):
"""
User has uploaded base64 encodings as JSON, this task needs to copy the data into
our internal binary format.
"""
log = logger.bind(pid=project_id, dp_id=dp_id)
log.info("Handling user provided base64 encodings")
new_child_span = lambda name: handle_raw_upload.tracer.start_active_span(name, child_of=handle_raw_upload.span)
with DBConn() as db:
if not check_project_exists(db, project_id):
log.info("Project deleted, stopping immediately")
return
# Get number of blocks + total number of encodings from database
expected_count, block_count = get_encoding_metadata(db, dp_id)
log.info(f"Expecting to handle {expected_count} encodings in {block_count} blocks")
mc = connect_to_object_store()
input_filename = Config.RAW_FILENAME_FMT.format(receipt_token)
raw_data = mc.get_object(Config.MINIO_BUCKET, input_filename)
with new_child_span('upload-encodings-to-db'):
# stream encodings with block ids from uploaded file
# convert each encoding to our internal binary format
# output into database for each block (temp or direct to minio?)
encoding_size, pipeline = convert_encodings_from_base64_to_binary(stream_json_clksnblocks(raw_data))
log.info(f"Starting pipeline to store {encoding_size}B sized encodings in database")
with DBConn() as db:
store_encodings_in_db(db, dp_id, pipeline, encoding_size)
log.info(f"Converted uploaded encodings of size {fmt_bytes(encoding_size)} into internal binary format. Number of blocks: {block_count}")
# As this is the first time we've seen the encoding size actually uploaded from this data provider
# We check it complies with the project encoding size.
try:
check_dataproviders_encoding(project_id, encoding_size)
except InvalidEncodingError as e:
log.warning(e.args[0])
handle_invalid_encoding_data(project_id, dp_id)
with DBConn() as conn:
with new_child_span('save-encoding-metadata'):
# Save the encoding size as metadata for this data provider
update_encoding_metadata_set_encoding_size(conn, dp_id, encoding_size)
update_encoding_metadata(conn, None, dp_id, 'ready')
# Now work out if all parties have added their data
if clks_uploaded_to_project(project_id, check_data_ready=True):
log.info("All parties' data present. Scheduling any queued runs")
check_for_executable_runs.delay(project_id, handle_raw_upload.get_serialized_span())
def handle_raw_upload(project_id, dp_id, receipt_token, parent_span=None):
# User has uploaded base64 encodings as JSON
log = logger.bind(pid=project_id, dp_id=dp_id)
log.info("Handling user provided base64 encodings")
with DBConn() as db:
if not check_project_exists(db, project_id):
log.info("Project deleted, stopping immediately")
return
expected_count = get_number_of_hashes(db, dp_id)
log.info(f"Expecting to handle {expected_count} encodings")
mc = connect_to_object_store()
# Input file is line separated base64 record encodings.
raw_file = Config.RAW_FILENAME_FMT.format(receipt_token)
raw_data_response = mc.get_object(Config.MINIO_BUCKET, raw_file)
# Set up streaming processing pipeline
buffered_stream = iterable_to_stream(raw_data_response.stream())
text_stream = io.TextIOWrapper(buffered_stream, newline='\n')
clkcounts = []
def filter_generator():
log.debug("Deserializing json filters")
first_encoding_size = None
for i, line in enumerate(text_stream):
ba = deserialize_bitarray(line)
yield (ba, i, ba.count())
clkcounts.append(ba.count())
encsize = len(ba)
if i == 0:
first_encoding_size = encsize
if encsize != first_encoding_size:
raise ValueError("Encodings were not all the same size")
log.info(f"Processed {len(clkcounts)} hashes")
# We peek at the first element as we need the encoding size
# for the ret of our processing pipeline
python_filters = more_itertools.peekable(filter_generator())
# Note the len of a bitarray is returned in bits but we require
# this to be a multiple of 8 so we use bytes.
uploaded_encoding_size = len(python_filters.peek()[0]) // 8
# This is the first time we've seen the encoding size from this data provider
try:
check_dataproviders_encoding(project_id, uploaded_encoding_size)
except InvalidEncodingError as e:
log.warning(e.args[0])
handle_invalid_encoding_data(project_id, dp_id)
with DBConn() as db:
# Save the encoding size as metadata
update_encoding_metadata_set_encoding_size(db, dp_id,
uploaded_encoding_size)
# Output file is our custom binary packed file
filename = Config.BIN_FILENAME_FMT.format(receipt_token)
bit_packed_element_size = binary_format(uploaded_encoding_size).size
num_bytes = expected_count * bit_packed_element_size
# If small enough preload the data into our redis cache
if expected_count < Config.ENTITY_CACHE_THRESHOLD:
log.info("Caching pickled clk data")
python_filters = list(python_filters)
cache.set_deserialized_filter(dp_id, python_filters)
else:
log.info("Not caching clk data as it is too large")
packed_filters = binary_pack_filters(python_filters,
uploaded_encoding_size)
packed_filter_stream = iterable_to_stream(packed_filters)
# Upload to object store
log.info(
f"Uploading {expected_count} encodings of size {uploaded_encoding_size} "
+ f"to object store. Total Size: {fmt_bytes(num_bytes)}")
mc.put_object(Config.MINIO_BUCKET,
filename,
data=packed_filter_stream,
length=num_bytes)
with DBConn() as conn:
update_encoding_metadata(conn, filename, dp_id, 'ready')
# Now work out if all parties have added their data
if clks_uploaded_to_project(project_id, check_data_ready=True):
log.info("All parties' data present. Scheduling any queued runs")
check_for_executable_runs.delay(
project_id, handle_raw_upload.get_serialized_span())
def pull_external_data(project_id, dp_id,
encoding_object_info,
blocks_object_info,
receipt_token, parent_span=None):
"""
Load encoding and blocking data from object store.
- pull blocking map into memory, create blocks in db
- stream encodings into DB and add encoding + blocks from in memory dict.
:param project_id: identifier for the project
:param dp_id:
:param encoding_object_info: a dictionary contains bucket and path of uploaded encoding
:param blocks_object_info: a dictionary contains bucket and path of uploaded blocks
:param receipt_token: token used to insert into database
"""
env_credentials = parse_minio_credentials({
'AccessKeyId': config.MINIO_ACCESS_KEY,
'SecretAccessKey': config.MINIO_SECRET_KEY
})
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(env_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)['blocks']
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, env_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 of size {size}")
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):
insert_encoding_metadata(conn, None, dp_id, receipt_token, encoding_count=count, block_count=block_count)
update_encoding_metadata_set_encoding_size(conn, dp_id, size)
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 ijson_encoding_iterator(encoding_stream):
binary_formatter = binary_format(size)
for encoding_id, encoding in zip(range(count), encoding_stream):
yield (
str(encoding_id),
binary_formatter.pack(encoding_id, deserialize_bytes(encoding)),
encoding_to_block_map[str(encoding_id)]
)
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)]
)
if object_name.endswith('.json'):
encodings_stream = ijson.items(io.BytesIO(encodings_stream.data), 'clks.item')
encoding_generator = ijson_encoding_iterator(encodings_stream)
else:
encoding_generator = encoding_iterator(encodings_stream)
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_generator, 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')
# # Now work out if all parties have added their data
if clks_uploaded_to_project(project_id):
logger.info("All parties data present. Scheduling any queued runs")
check_for_executable_runs.delay(project_id, serialize_span(parent_span))
def project_clks_post(project_id):
"""
Update a project to provide encoded PII data.
"""
headers = request.headers
log, parent_span = bind_log_and_span(project_id)
log.debug("Starting data upload request")
token = precheck_upload_token(project_id, headers, parent_span)
receipt_token = generate_code()
with DBConn() as conn:
dp_id = db.get_dataprovider_id(conn, token)
project_encoding_size = db.get_project_schema_encoding_size(
conn, project_id)
upload_state_updated = db.is_dataprovider_allowed_to_upload_and_lock(
conn, dp_id)
# get flag use_blocking from table projects
uses_blocking = get_project_column(conn, project_id, 'uses_blocking')
if not upload_state_updated:
return safe_fail_request(
403, "This token has already been used to upload clks.")
log = log.bind(dp_id=dp_id)
log.info("Receiving CLK data.")
with opentracing.tracer.start_span('upload-clk-data',
child_of=parent_span) as span:
span.set_tag("project_id", project_id)
try:
if headers['Content-Type'] == "application/json":
span.set_tag("content-type", 'json')
# TODO: Previously, we were accessing the CLKs in a streaming fashion to avoid parsing the json in one hit. This
# enables running the web frontend with less memory.
# However, as connexion is very, very strict about input validation when it comes to json, it will always
# consume the stream first to validate it against the spec. Thus the backflip to fully reading the CLks as
# json into memory. -> issue #184
handle_encoding_upload_json(project_id,
dp_id,
get_json(),
receipt_token,
uses_blocking,
parent_span=span)
log.info("Job scheduled to handle users upload")
elif headers['Content-Type'] == "application/octet-stream":
span.set_tag("content-type", 'binary')
log.info("Handling binary CLK upload")
try:
count, size = check_binary_upload_headers(headers)
log.info(
f"Headers tell us to expect {count} encodings of {size} bytes"
)
span.log_kv({'count': count, 'size': size})
except Exception:
log.warning(
"Upload failed due to problem with headers in binary upload"
)
raise
# Check against project level encoding size (if it has been set)
if project_encoding_size is not None and size != project_encoding_size:
# fail fast - we haven't stored the encoded data yet
return safe_fail_request(
400,
"Upload 'Hash-Size' doesn't match project settings")
# TODO actually stream the upload data straight to Minio. Currently we can't because
# connexion has already read the data before our handler is called!
# https://github.com/zalando/connexion/issues/592
# stream = get_stream()
stream = BytesIO(request.data)
expected_bytes = binary_format(size).size * count
log.debug(
f"Stream size is {len(request.data)} B, and we expect {expected_bytes} B"
)
if len(request.data) != expected_bytes:
safe_fail_request(
400,
"Uploaded data did not match the expected size. Check request headers are correct"
)
try:
upload_clk_data_binary(project_id, dp_id, stream,
receipt_token, count, size)
except ValueError:
safe_fail_request(
400,
"Uploaded data did not match the expected size. Check request headers are correct."
)
else:
safe_fail_request(400, "Content Type not supported")
except ProblemException as e:
# Have an exception that is safe for the user. We reset the upload state to
# allow the user to try upload again.
log.info(
f"Problem occurred, returning status={e.status} - {e.detail}")
with DBConn() as conn:
db.set_dataprovider_upload_state(conn,
dp_id,
state='not_started')
raise
except Exception as e:
log.warning("Unhandled error occurred during data upload")
log.exception(e)
with DBConn() as conn:
db.set_dataprovider_upload_state(conn, dp_id, state='error')
safe_fail_request(
500, "Sorry, the server couldn't handle that request")
with DBConn() as conn:
db.set_dataprovider_upload_state(conn, dp_id, state='done')
# Now work out if all parties have added their data
if clks_uploaded_to_project(project_id):
logger.info("All parties data present. Scheduling any queued runs")
check_for_executable_runs.delay(project_id,
serialize_span(parent_span))
return {'message': 'Updated', 'receipt_token': receipt_token}, 201
def project_binaryclks_post(project_id):
"""
Update a project to provide encoded PII data.
"""
log, parent_span = bind_log_and_span(project_id)
headers = request.headers
token = precheck_upload_token(project_id, headers, parent_span)
with DBConn() as conn:
dp_id = db.get_dataprovider_id(conn, token)
project_encoding_size = db.get_project_schema_encoding_size(
conn, project_id)
upload_state_updated = db.is_dataprovider_allowed_to_upload_and_lock(
conn, dp_id)
if not upload_state_updated:
return safe_fail_request(
403, "This token has already been used to upload clks.")
log = log.bind(dp_id=dp_id)
log.info("Receiving CLK data.")
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)
try:
if headers['Content-Type'] == "application/octet-stream":
span.set_tag("content-type", 'binary')
log.info("Handling binary CLK upload")
try:
count, size = check_binary_upload_headers(headers)
log.info(
f"Headers tell us to expect {count} encodings of {size} bytes"
)
span.log_kv({'count': count, 'size': size})
except Exception:
log.warning(
"Upload failed due to problem with headers in binary upload"
)
raise
# Check against project level encoding size (if it has been set)
if project_encoding_size is not None and size != project_encoding_size:
# fail fast - we haven't stored the encoded data yet
return safe_fail_request(
400,
"Upload 'Hash-Size' doesn't match project settings")
# TODO actually stream the upload data straight to Minio. Currently we can't because
# connexion has already read the data before our handler is called!
# https://github.com/zalando/connexion/issues/592
# stream = get_stream()
stream = BytesIO(request.data)
converted_stream = include_encoding_id_in_binary_stream(
stream, size, count)
expected_bytes = size * count
log.debug(
f"Stream size is {len(request.data)} B, and we expect {expected_bytes} B"
)
if len(request.data) != expected_bytes:
safe_fail_request(
400,
"Uploaded data did not match the expected size. Check request headers are correct"
)
try:
upload_clk_data_binary(project_id, dp_id, converted_stream,
receipt_token, count, size)
except ValueError:
safe_fail_request(
400,
"Uploaded data did not match the expected size. Check request headers are correct."
)
else:
safe_fail_request(400, "Content Type not supported")
except Exception:
log.warning(
"The dataprovider was not able to upload their clks,"
" re-enable the corresponding upload token to be used.")
with DBConn() as conn:
db.set_dataprovider_upload_state(conn, dp_id, state='error')
raise
with DBConn() as conn:
db.set_dataprovider_upload_state(conn, dp_id, state='done')
# Now work out if all parties have added their data
if clks_uploaded_to_project(project_id):
logger.info("All parties data present. Scheduling any queued runs")
check_for_executable_runs.delay(project_id,
serialize_span(parent_span))
return {'message': 'Updated', 'receipt_token': receipt_token}, 201