80 lines
2.5 KiB
Ruby
80 lines
2.5 KiB
Ruby
require 'json'
|
|
require_relative 'storage_plugin'
|
|
|
|
class DataPipeline
|
|
#
|
|
# options
|
|
# file: path of the file
|
|
# format: :array (one json array) or :lines (one object per line, ndjson)
|
|
# default from the extension: .ndjson / .jsonl -> :lines, else :array
|
|
# root: read only: path to the array inside the document, e.g. 'data.items'
|
|
# a single object (instead of an array) is read as one row
|
|
# pretty: write :array pretty printed (default false)
|
|
#
|
|
class JsonStorage < StoragePlugin
|
|
LINE_EXTENSIONS = %w[.ndjson .jsonl].freeze
|
|
|
|
def initialize **options, &block
|
|
@options = options
|
|
@file_name = options.fetch(:file) { raise ArgumentError, 'JsonStorage: option file missing' }
|
|
@format = options.fetch(:format) { LINE_EXTENSIONS.include?(File.extname(@file_name)) ? :lines : :array }
|
|
raise ArgumentError, "JsonStorage: format #{@format} unknown, use :array or :lines" unless %i[array lines].include?(@format)
|
|
|
|
yield if block_given?
|
|
end
|
|
|
|
# without block all rows, with block row by row (ndjson is read line by line)
|
|
def read(&block)
|
|
rows = @format == :lines ? read_lines : read_array
|
|
return rows.to_a unless block
|
|
|
|
rows.each(&block)
|
|
nil
|
|
end
|
|
|
|
def write rows
|
|
if @format == :lines
|
|
File.open(@file_name, 'w') { |f| rows.each { |row| f.puts JSON.generate(row) } }
|
|
else
|
|
File.write(@file_name, @options[:pretty] ? JSON.pretty_generate(rows) : JSON.generate(rows))
|
|
end
|
|
end
|
|
|
|
# streaming: ndjson line by line, array as "[row,row,...]" closed in close
|
|
def write_row row
|
|
@file ||= File.open(@file_name, 'w').tap { |f| f.write('[') if @format == :array }
|
|
if @format == :array
|
|
@file.write(',') if @rows_written
|
|
@file.write(JSON.generate(row))
|
|
else
|
|
@file.puts JSON.generate(row)
|
|
end
|
|
@rows_written = true
|
|
end
|
|
|
|
def close
|
|
return unless @file
|
|
|
|
@file.write(']') if @format == :array
|
|
@file.close
|
|
@file = nil
|
|
@rows_written = false
|
|
end
|
|
|
|
private
|
|
|
|
def read_lines
|
|
File.foreach(@file_name).lazy.map(&:strip).reject(&:empty?).map { |line| JSON.parse(line) }
|
|
end
|
|
|
|
def read_array
|
|
data = JSON.parse(File.read(@file_name))
|
|
@options[:root].to_s.split('.').each { |key| data = data.fetch(key) } if @options[:root]
|
|
data = [data] if data.is_a?(Hash) # a single object is one row
|
|
raise ArgumentError, "JsonStorage: #{@file_name} does not contain an array" unless data.is_a?(Array)
|
|
|
|
data
|
|
end
|
|
end
|
|
end
|