projects.mako.cc - selectricity/blob - vendor/plugins/yaml_db/lib/yaml_db.rb

   1 require 'rubygems'
   2 require 'yaml'
   3 require 'active_record'
   4
   5
   6 module YamlDb
   7         def self.dump(filename)
   8                 disable_logger
   9                 YamlDb::Dump.dump(File.new(filename, "w"))
  10                 reenable_logger
  11         end
  12
  13         def self.load(filename)
  14                 disable_logger
  15                 YamlDb::Load.load(File.new(filename, "r"))
  16                 reenable_logger
  17         end
  18
  19         def self.disable_logger
  20                 @@old_logger = ActiveRecord::Base.logger
  21                 ActiveRecord::Base.logger = nil
  22         end
  23
  24         def self.reenable_logger
  25                 ActiveRecord::Base.logger = @@old_logger
  26         end
  27 end
  28
  29
  30 module YamlDb::Utils
  31         def self.chunk_records(records)
  32                 yaml = [ records ].to_yaml
  33                 yaml.sub!("--- \n", "")
  34                 yaml.sub!('- - -', '  - -')
  35                 yaml
  36         end
  37
  38         def self.unhash(hash, keys)
  39                 keys.map { |key| hash[key] }
  40         end
  41
  42         def self.unhash_records(records, keys)
  43                 records.each_with_index do |record, index|
  44                         records[index] = unhash(record, keys)
  45                 end
  46
  47                 records
  48         end
  49
  50         def self.convert_booleans(records, columns)
  51                 records.each do |record|
  52                         columns.each do |column|
  53                                 next if is_boolean(record[column])
  54                                 record[column] = (record[column] == 't' or record[column] == '1')
  55                         end
  56                 end
  57                 records
  58         end
  59
  60         def self.boolean_columns(table)
  61                 columns = ActiveRecord::Base.connection.columns(table).reject { |c| c.type != :boolean }
  62                 columns.map { |c| c.name }
  63         end
  64
  65         def self.is_boolean(value)
  66                 value.kind_of?(TrueClass) or value.kind_of?(FalseClass)
  67         end
  68
  69         def self.quote_table(table)
  70                 ActiveRecord::Base.connection.quote_table_name(table)
  71         end
  72 end
  73
  74
  75 module YamlDb::Dump
  76         def self.dump(io)
  77                 tables.each do |table|
  78                         dump_table(io, table)
  79                 end
  80         end
  81
  82         def self.tables
  83                 ActiveRecord::Base.connection.tables.reject { |table| ['schema_info', 'schema_migrations'].include?(table) }
  84         end
  85
  86         def self.dump_table(io, table)
  87                 return if table_record_count(table).zero?
  88
  89                 dump_table_columns(io, table)
  90                 dump_table_records(io, table)
  91         end
  92
  93         def self.dump_table_columns(io, table)
  94                 io.write("\n")
  95                 io.write({ table => { 'columns' => table_column_names(table) } }.to_yaml)
  96         end
  97
  98         def self.dump_table_records(io, table)
  99                 table_record_header(io)
 100
 101                 column_names = table_column_names(table)
 102
 103                 each_table_page(table) do |records|
 104                         rows = YamlDb::Utils.unhash_records(records, column_names)
 105                         io.write(YamlDb::Utils.chunk_records(records))
 106                 end
 107         end
 108
 109         def self.table_record_header(io)
 110                 io.write("  records: \n")
 111         end
 112
 113         def self.table_column_names(table)
 114                 ActiveRecord::Base.connection.columns(table).map { |c| c.name }
 115         end
 116
 117         def self.each_table_page(table, records_per_page=1000)
 118                 total_count = table_record_count(table)
 119                 pages = (total_count.to_f / records_per_page).ceil - 1
 120                 id = table_column_names(table).first
 121                 boolean_columns = YamlDb::Utils.boolean_columns(table)
 122                 quoted_table_name = YamlDb::Utils.quote_table(table)
 123
 124                 (0..pages).to_a.each do |page|
 125                         sql = ActiveRecord::Base.connection.add_limit_offset!("SELECT * FROM #{quoted_table_name} ORDER BY #{id}",
 126                                 :limit => records_per_page, :offset => records_per_page * page
 127                         )
 128                         records = ActiveRecord::Base.connection.select_all(sql)
 129                         records = YamlDb::Utils.convert_booleans(records, boolean_columns)
 130                         yield records
 131                 end
 132         end
 133
 134         def self.table_record_count(table)
 135                 ActiveRecord::Base.connection.select_one("SELECT COUNT(*) FROM #{YamlDb::Utils.quote_table(table)}").values.first.to_i
 136         end
 137 end
 138
 139
 140 module YamlDb::Load
 141         def self.load(io)
 142                 ActiveRecord::Base.connection.transaction do
 143                         YAML.load_documents(io) do |ydoc|
 144                                 ydoc.keys.each do |table_name|
 145                                         next if ydoc[table_name].nil?
 146                                         load_table(table_name, ydoc[table_name])
 147                                 end
 148                         end
 149                 end
 150         end
 151
 152         def self.truncate_table(table)
 153                 begin
 154                         ActiveRecord::Base.connection.execute("TRUNCATE #{YamlDb::Utils.quote_table(table)}")
 155                 rescue Exception
 156                         ActiveRecord::Base.connection.execute("DELETE FROM #{YamlDb::Utils.quote_table(table)}")
 157                 end
 158         end
 159
 160         def self.load_table(table, data)
 161                 column_names = data['columns']
 162                 truncate_table(table)
 163                 load_records(table, column_names, data['records'])
 164                 reset_pk_sequence!(table)
 165         end
 166
 167         def self.load_records(table, column_names, records)
 168                 quoted_column_names = column_names.map { |column| ActiveRecord::Base.connection.quote_column_name(column) }.join(',')
 169                 quoted_table_name = YamlDb::Utils.quote_table(table)
 170                 records.each do |record|
 171                         ActiveRecord::Base.connection.execute("INSERT INTO #{quoted_table_name} (#{quoted_column_names}) VALUES (#{record.map { |r| ActiveRecord::Base.connection.quote(r) }.join(',')})")
 172                 end
 173         end
 174
 175         def self.reset_pk_sequence!(table_name)
 176                 if ActiveRecord::Base.connection.respond_to?(:reset_pk_sequence!)
 177                         ActiveRecord::Base.connection.reset_pk_sequence!(table_name)
 178                 end
 179         end
 180 end