def _write_to_entity(*, entity: EntityType) -> RespTuple:
from snuba.processor import InsertBatch
rows: MutableSequence[WriterTableRow] = []
offset_base = int(round(time.time() * 1000))
writable_storage = entity.get_writable_storage()
assert writable_storage is not None
table_writer = writable_storage.get_table_writer()
for index, message in enumerate(json.loads(http_request.data)):
offset = offset_base + index
processed_message = (
table_writer.get_stream_loader()
.get_processor()
.process_message(
message,
KafkaMessageMetadata(
offset=offset, partition=0, timestamp=datetime.utcnow()
),
)
)
if processed_message:
assert isinstance(processed_message, InsertBatch)
rows.extend(processed_message.rows)
BatchWriterEncoderWrapper(
table_writer.get_batch_writer(metrics),
JSONRowEncoder(),
).write(rows)
return ("ok", 200, {"Content-Type": "text/plain"})
def enforce_table_writer(entity: Entity) -> TableWriter:
writable_storage = entity.get_writable_storage()
assert (
writable_storage is not None
), f"Entity {ENTITY_NAME_LOOKUP[entity]} does not have a writable storage."
return writable_storage.get_table_writer()