Skip to content
This repository
Browse code

Merge pull request #5698 from dougcole/support_postgresql_partitioning

Support postgresql partitioning by making INSERT RETURNING optional
  • Loading branch information...
commit 8de4d71f5dab243f2c66e1695ccfabc0bbf98c9b 2 parents 3981a68 + cd6ddc8
Aaron Patterson tenderlove authored
4 activerecord/lib/active_record/connection_adapters/abstract/database_statements.rb
@@ -59,7 +59,7 @@ def exec_query(sql, name = 'SQL', binds = [])
59 59 # Executes insert +sql+ statement in the context of this connection using
60 60 # +binds+ as the bind substitutes. +name+ is logged along with
61 61 # the executed +sql+ statement.
62   - def exec_insert(sql, name, binds)
  62 + def exec_insert(sql, name, binds, pk = nil, sequence_name = nil)
63 63 exec_query(sql, name, binds)
64 64 end
65 65
@@ -87,7 +87,7 @@ def exec_update(sql, name, binds)
87 87 # passed in as +id_value+.
88 88 def insert(arel, name = nil, pk = nil, id_value = nil, sequence_name = nil, binds = [])
89 89 sql, binds = sql_for_insert(to_sql(arel, binds), pk, id_value, sequence_name, binds)
90   - value = exec_insert(sql, name, binds)
  90 + value = exec_insert(sql, name, binds, pk, sequence_name)
91 91 id_value || last_inserted_id(value)
92 92 end
93 93
2  activerecord/lib/active_record/connection_adapters/mysql2_adapter.rb
@@ -226,7 +226,7 @@ def insert_sql(sql, name = nil, pk = nil, id_value = nil, sequence_name = nil)
226 226 end
227 227 alias :create :insert_sql
228 228
229   - def exec_insert(sql, name, binds)
  229 + def exec_insert(sql, name, binds, pk = nil, sequence_name = nil)
230 230 execute to_sql(sql, binds), name
231 231 end
232 232
47 activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb
@@ -17,7 +17,7 @@ def postgresql_connection(config) # :nodoc:
17 17 # Forward any unused config params to PGconn.connect.
18 18 [:statement_limit, :encoding, :min_messages, :schema_search_path,
19 19 :schema_order, :adapter, :pool, :wait_timeout, :template,
20   - :reaping_frequency].each do |key|
  20 + :reaping_frequency, :insert_returning].each do |key|
21 21 conn_params.delete key
22 22 end
23 23 conn_params.delete_if { |k,v| v.nil? }
@@ -258,6 +258,8 @@ def simplified_type(field_type)
258 258 # <encoding></tt> call on the connection.
259 259 # * <tt>:min_messages</tt> - An optional client min messages that is used in a
260 260 # <tt>SET client_min_messages TO <min_messages></tt> call on the connection.
  261 + # * <tt>:insert_returning</tt> - An optional boolean to control the use or <tt>RETURNING</tt> for <tt>INSERT<tt> statements
  262 + # defaults to true.
261 263 #
262 264 # Any further options are used as connection parameters to libpq. See
263 265 # http://www.postgresql.org/docs/9.1/static/libpq-connect.html for the
@@ -405,6 +407,7 @@ def initialize(connection, logger, connection_parameters, config)
405 407
406 408 initialize_type_map
407 409 @local_tz = execute('SHOW TIME ZONE', 'SCHEMA').first["TimeZone"]
  410 + @use_insert_returning = @config.key?(:insert_returning) ? @config[:insert_returning] : true
408 411 end
409 412
410 413 # Clears the prepared statements cache.
@@ -666,8 +669,11 @@ def insert_sql(sql, name = nil, pk = nil, id_value = nil, sequence_name = nil)
666 669 pk = primary_key(table_ref) if table_ref
667 670 end
668 671
669   - if pk
  672 + if pk && use_insert_returning?
670 673 select_value("#{sql} RETURNING #{quote_column_name(pk)}")
  674 + elsif pk
  675 + super
  676 + last_insert_id_value(sequence_name || default_sequence_name(table_ref, pk))
671 677 else
672 678 super
673 679 end
@@ -782,11 +788,27 @@ def sql_for_insert(sql, pk, id_value, sequence_name, binds)
782 788 pk = primary_key(table_ref) if table_ref
783 789 end
784 790
785   - sql = "#{sql} RETURNING #{quote_column_name(pk)}" if pk
  791 + if pk && use_insert_returning?
  792 + sql = "#{sql} RETURNING #{quote_column_name(pk)}"
  793 + end
786 794
787 795 [sql, binds]
788 796 end
789 797
  798 + def exec_insert(sql, name, binds, pk = nil, sequence_name = nil)
  799 + val = exec_query(sql, name, binds)
  800 + if !use_insert_returning? && pk
  801 + unless sequence_name
  802 + table_ref = extract_table_ref_from_insert_sql(sql)
  803 + sequence_name = default_sequence_name(table_ref, pk)
  804 + return val unless sequence_name
  805 + end
  806 + last_insert_id_result(sequence_name)
  807 + else
  808 + val
  809 + end
  810 + end
  811 +
790 812 # Executes an UPDATE query and returns the number of affected tuples.
791 813 def update_sql(sql, name = nil)
792 814 super.cmd_tuples
@@ -1027,7 +1049,9 @@ def client_min_messages=(level)
1027 1049
1028 1050 # Returns the sequence name for a table's primary key or some other specified key.
1029 1051 def default_sequence_name(table_name, pk = nil) #:nodoc:
1030   - serial_sequence(table_name, pk || 'id').split('.').last
  1052 + result = serial_sequence(table_name, pk || 'id')
  1053 + return nil unless result
  1054 + result.split('.').last
1031 1055 rescue ActiveRecord::StatementInvalid
1032 1056 "#{table_name}_#{pk || 'id'}_seq"
1033 1057 end
@@ -1235,6 +1259,10 @@ def extract_schema_and_table(name)
1235 1259 end
1236 1260 end
1237 1261
  1262 + def use_insert_returning?
  1263 + @use_insert_returning
  1264 + end
  1265 +
1238 1266 protected
1239 1267 # Returns the version of the connected PostgreSQL server.
1240 1268 def postgresql_version
@@ -1364,8 +1392,15 @@ def configure_connection
1364 1392
1365 1393 # Returns the current ID of a table's sequence.
1366 1394 def last_insert_id(sequence_name) #:nodoc:
1367   - r = exec_query("SELECT currval($1)", 'SQL', [[nil, sequence_name]])
1368   - Integer(r.rows.first.first)
  1395 + Integer(last_insert_id_value(sequence_name))
  1396 + end
  1397 +
  1398 + def last_insert_id_value(sequence_name)
  1399 + last_insert_id_result(sequence_name).rows.first.first
  1400 + end
  1401 +
  1402 + def last_insert_id_result(sequence_name) #:nodoc:
  1403 + exec_query("SELECT currval($1)", 'SQL', [[nil, sequence_name]])
1369 1404 end
1370 1405
1371 1406 # Executes a SELECT query and returns the results, performing any data type
31 activerecord/test/cases/adapters/postgresql/postgresql_adapter_test.rb
@@ -49,6 +49,33 @@ def test_insert_sql_with_no_space_after_table_name
49 49 assert_equal expect, id
50 50 end
51 51
  52 + def test_insert_sql_with_returning_disabled
  53 + connection = connection_without_insert_returning
  54 + id = connection.insert_sql("insert into postgresql_partitioned_table_parent (number) VALUES (1)")
  55 + expect = connection.query('select max(id) from postgresql_partitioned_table_parent').first.first
  56 + assert_equal expect, id
  57 + end
  58 +
  59 + def test_exec_insert_with_returning_disabled
  60 + connection = connection_without_insert_returning
  61 + result = connection.exec_insert("insert into postgresql_partitioned_table_parent (number) VALUES (1)", nil, [], 'id', 'postgresql_partitioned_table_parent_id_seq')
  62 + expect = connection.query('select max(id) from postgresql_partitioned_table_parent').first.first
  63 + assert_equal expect, result.rows.first.first
  64 + end
  65 +
  66 + def test_exec_insert_with_returning_disabled_and_no_sequence_name_given
  67 + connection = connection_without_insert_returning
  68 + result = connection.exec_insert("insert into postgresql_partitioned_table_parent (number) VALUES (1)", nil, [], 'id')
  69 + expect = connection.query('select max(id) from postgresql_partitioned_table_parent').first.first
  70 + assert_equal expect, result.rows.first.first
  71 + end
  72 +
  73 + def test_sql_for_insert_with_returning_disabled
  74 + connection = connection_without_insert_returning
  75 + result = connection.sql_for_insert('sql', nil, nil, nil, 'binds')
  76 + assert_equal ['sql', 'binds'], result
  77 + end
  78 +
52 79 def test_serial_sequence
53 80 assert_equal 'public.accounts_id_seq',
54 81 @connection.serial_sequence('accounts', 'id')
@@ -204,6 +231,10 @@ def insert(ctx, data)
204 231
205 232 ctx.exec_insert(sql, 'SQL', binds)
206 233 end
  234 +
  235 + def connection_without_insert_returning
  236 + ActiveRecord::Base.postgresql_connection(ActiveRecord::Base.configurations['arunit'].merge(:insert_returning => false))
  237 + end
207 238 end
208 239 end
209 240 end
29 activerecord/test/schema/postgresql_specific_schema.rb
... ... @@ -1,8 +1,8 @@
1 1 ActiveRecord::Schema.define do
2 2
3 3 %w(postgresql_tsvectors postgresql_hstores postgresql_arrays postgresql_moneys postgresql_numbers postgresql_times postgresql_network_addresses postgresql_bit_strings
4   - postgresql_oids postgresql_xml_data_type defaults geometrics postgresql_timestamp_with_zones).each do |table_name|
5   - execute "DROP TABLE IF EXISTS #{quote_table_name table_name}"
  4 + postgresql_oids postgresql_xml_data_type defaults geometrics postgresql_timestamp_with_zones postgresql_partitioned_table postgresql_partitioned_table_parent).each do |table_name|
  5 + execute "DROP TABLE IF EXISTS #{quote_table_name table_name}"
6 6 end
7 7
8 8 execute 'DROP SEQUENCE IF EXISTS companies_nonstd_seq CASCADE'
@@ -10,6 +10,8 @@
10 10 execute "ALTER TABLE companies ALTER COLUMN id SET DEFAULT nextval('companies_nonstd_seq')"
11 11 execute 'DROP SEQUENCE IF EXISTS companies_id_seq'
12 12
  13 + execute 'DROP FUNCTION IF EXISTS partitioned_insert_trigger()'
  14 +
13 15 %w(accounts_id_seq developers_id_seq projects_id_seq topics_id_seq customers_id_seq orders_id_seq).each do |seq_name|
14 16 execute "SELECT setval('#{seq_name}', 100)"
15 17 end
@@ -125,6 +127,29 @@
125 127 );
126 128 _SQL
127 129
  130 + execute <<_SQL
  131 + CREATE TABLE postgresql_partitioned_table_parent (
  132 + id SERIAL PRIMARY KEY,
  133 + number integer
  134 + );
  135 + CREATE TABLE postgresql_partitioned_table ( )
  136 + INHERITS (postgresql_partitioned_table_parent);
  137 +
  138 + CREATE OR REPLACE FUNCTION partitioned_insert_trigger()
  139 + RETURNS TRIGGER AS $$
  140 + BEGIN
  141 + INSERT INTO postgresql_partitioned_table VALUES (NEW.*);
  142 + RETURN NULL;
  143 + END;
  144 + $$
  145 + LANGUAGE plpgsql;
  146 +
  147 + CREATE TRIGGER insert_partitioning_trigger
  148 + BEFORE INSERT ON postgresql_partitioned_table_parent
  149 + FOR EACH ROW EXECUTE PROCEDURE partitioned_insert_trigger();
  150 +_SQL
  151 +
  152 +
128 153 begin
129 154 execute <<_SQL
130 155 CREATE TABLE postgresql_xml_data_type (

0 comments on commit 8de4d71

Please sign in to comment.
Something went wrong with that request. Please try again.