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()