open-uriライブラリ10kbの罠

RakeタスクでとあるCSVファイルをダウンロードしながら(CSV.foreach)、データの加工やDBへの保存処理を行うスクリプトを作成していたところ、以下にあるような課題に当たって、ヘルパーメソッドを作成したので背景とともに残しておく

背景と課題

1. ローカルで手軽にRakeタスクを実行したいが、HerokuでホストされているDBへのIOに時間がかかる
   • heroku runコマンドで実行するRakeタスクの引数に、GCSの公開バケットのURLを指定して、DBとスクリプトの実行環境を近づけることに
2. GCSのバケットにあらかじめアップロードしておいたCSVファイルをURI.parse('<https://storage.google.xxx.com/sample.csv>').openを使用し、http経由でダウンロードするようにしたところ、10kb以下のファイルがStringIOで返ってくるためうまくいかない

openメソッドはOpenURI::Buffer のStringMax 定数(10kb)を見てバッファ(StringIO)で処理するかどうかを判断しているおり、ファイル中のデータサイズが10kb以下の場合でそれぞれ、返されるオブジェクトの種類が異なる。 一方、CSV.foreach はIO stream可能なオブジェクトを引数にする必要があるため、一工夫必要

つくったもの

• CSV.foreachで扱えるIO stream可能なオブジェクトを返すヘルパーメソッドを作成してみた
• ついでにパス・URL指定どちらも指定できるように

require 'open-uri'

def read_as_path_or_io(path_or_url)
  safe_uri = URI.parse(path_or_url)

  # localのファイルシステムから参照
  if safe_uri.scheme.nil?
    return path_or_url
  end

  stream = safe_uri.open

  # 10kb以上のファイルをhttp or ftpから参照
  return stream if stream.respond_to?(:path)

  # 10kb以下のファイルをhttp or ftpから参照
  Tempfile.new.tap do |file|
    file.binmode
    IO.copy_stream(stream, file)
    stream.close
    file.rewind
  end
end

# 呼び出し
temp_file = read_as_path_or_io('ローカルのファイルパス or URL')
CSV.foreach(temp_file, encoding: 'UTF-8', headers: true).do |csv|
  p csv
end

参考

https://github.com/ruby/csv/blob/master/lib/csv.rb https://github.com/gitlabhq/gitlabhq/blob/master/spec/support/helpers/rake_helpers.rb https://stackoverflow.com/questions/31527154/rails-open-returns-stringio-instead-of-tempfile https://docs.ruby-lang.org/ja/latest/library/open=2duri.html https://stackoverflow.com/questions/31527154/rails-open-returns-stringio-instead-of-tempfile https://docs.ruby-lang.org/ja/latest/class/Tempfile.html