]> projects.mako.cc - selectricity-live/blob - vendor/plugins/yaml_db/lib/yaml_db.rb
Fix errors when viewing results as election creator before election close.
[selectricity-live] / 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

Benjamin Mako Hill || Want to submit a patch?