Function get_buffer_writer

trinity/buffer/buffer.py:26–45

Get a buffer writer for the given dataset name.

(config: BufferStorageConfig)

def get_buffer_writer(config: BufferStorageConfig) -> BufferWriter:
    """Get a buffer writer for the given dataset name."""
    if not isinstance(config, StorageConfig):
        storage_config: StorageConfig = config.to_storage_config()
    else:
        storage_config = config
    if storage_config.storage_type == StorageType.SQL.value:
        from trinity.buffer.writer.sql_writer import SQLWriter

        return SQLWriter(storage_config)
    elif storage_config.storage_type == StorageType.QUEUE.value:
        from trinity.buffer.writer.queue_writer import QueueWriter

        return QueueWriter(storage_config)
    elif storage_config.storage_type == StorageType.FILE.value:
        from trinity.buffer.writer.file_writer import JSONWriter

        return JSONWriter(storage_config)
    else:
        raise ValueError(f"{storage_config.storage_type} not supported.")