lib/gtfs_loader.rb

   1 # frozen_string_literal: true
   2
   3 require "csv"
   4 require "pry"
   5 require "sqlite3"
   6
   7 class GtfsLoader
   8   def initialize
   9     @db = nil
  10   end
  11
  12   def create_tables
  13     calendar_sql = <<~EOS
  14       CREATE TABLE calendar (
  15         service_id VARCHAR(26),
  16         monday CHAR(1),
  17         tuesday CHAR(1),
  18         wednesday CHAR(1),
  19         thursday CHAR(1),
  20         friday CHAR(1),
  21         saturday CHAR(1),
  22         sunday CHAR(1),
  23         start_date CHAR(8),
  24         end_date CHAR(8)
  25       );
  26     EOS
  27     calendar_dates_sql = <<~EOS
  28       CREATE TABLE calendar_dates (
  29         service_id VARCHAR(26),
  30         calendar_date CHAR(8),
  31         exception_type CHAR(1)
  32       );
  33     EOS
  34     routes_sql = <<~EOS
  35       CREATE TABLE routes (
  36         route_id VARCHAR(8),
  37         route_short_name VARCHAR(16),
  38         route_type CHAR(1),
  39         route_color VARCHAR(6),
  40         route_text_color VARCHAR(6)
  41       );
  42     EOS
  43     stop_times_sql = <<~EOS
  44       CREATE TABLE stop_times (
  45         trip_id VARCHAR(35),
  46         arrival_time VARCHAR(8),
  47         departure_time VARCHAR(8),
  48         stop_id VARCHAR(5),
  49         stop_sequence VARCHAR(2),
  50         pickup_type CHAR(1),
  51         drop_off_type CHAR(1)
  52       );
  53     EOS
  54     stops_sql = <<~EOS
  55       CREATE TABLE stops (
  56         stop_id VARCHAR(5),
  57         stop_code VARCHAR(4),
  58         stop_name VARCHAR(48),
  59         stop_lat FLOAT,
  60         stop_lon FLOAT
  61       );
  62     EOS
  63     trips_sql = <<~EOS
  64       CREATE TABLE trips (
  65         route_id VARCHAR(8),
  66         service_id VARCHAR(26),
  67         trip_id VARCHAR(35),
  68         trip_headsign VARCHAR(48),
  69         direction_id CHAR(1),
  70         block_id VARCHAR(8),
  71         shape_id VARCHAR(7)
  72       );
  73     EOS
  74
  75     [
  76       calendar_sql,
  77       calendar_dates_sql,
  78       routes_sql,
  79       stop_times_sql,
  80       stops_sql,
  81       trips_sql,
  82     ].each do |sql|
  83       db.execute sql
  84     end
  85   end
  86
  87   TABLE_COLUMNS = {
  88     calendar: %w(service_id monday tuesday wednesday thursday friday saturday sunday start_date end_date),
  89     calendar_dates: %w(service_id date exception_type),
  90     routes: %w(route_id route_short_name route_type route_color route_text_color),
  91     stop_times: %w(trip_id arrival_time departure_time stop_id stop_sequence pickup_type drop_off_type),
  92     stops: %w(stop_id stop_code stop_name stop_lat stop_lon),
  93     trips: %w(route_id service_id trip_id trip_headsign direction_id block_id shape_id),
  94   }
  95
  96   def load_table(base_path, table_name)
  97     print "Loading #{table_name}..."
  98     $stdout.flush
  99     db.execute("BEGIN TRANSACTION;")
 100
 101     csv_cols = TABLE_COLUMNS[table_name.to_sym]
 102     sql_cols = csv_cols.map{|name| "date" == name ? "calendar_date" : name}
 103     count = 0
 104
 105     CSV.foreach("#{base_path}/#{table_name}.txt", headers: true, col_sep: ",") do |row|
 106       count += 1
 107       values = csv_cols.map{|col| row[col.to_s]}
 108       placeholders = (["?"] * values.count).join(",")
 109       sql = "INSERT INTO #{table_name} (#{sql_cols.join(",")}) VALUES (#{placeholders});"
 110
 111       db.execute(sql, values)
 112
 113       if (0 == count % 10_000)
 114         db.execute("COMMIT;")
 115         db.execute("BEGIN TRANSACTION;")
 116         print "."
 117         $stdout.flush
 118       end
 119     end
 120
 121     db.execute("COMMIT;")
 122     puts " done."
 123   end
 124
 125   def load_tables
 126     %w(calendar_dates calendar routes stops stop_times trips).each do |table|
 127       load_table("gtfs/2022-12-14/", table)
 128     end
 129   end
 130
 131   private
 132
 133   def db
 134     @db ||= SQLite3::Database.new("gtfs.db")
 135   end
 136 end