[WIP] Use RETURNING statement for INSERT statements in Oracle

lib/sequel/adapters/oracle.rb

@@ -13,23 +13,27 @@ class Database < Sequel::Database
       # ORA-01012: not logged on
       # ORA-03113: end-of-file on communication channel
       # ORA-03114: not connected to ORACLE
-      CONNECTION_ERROR_CODES = [ 28, 1012, 3113, 3114 ]      
-      
+      CONNECTION_ERROR_CODES = [ 28, 1012, 3113, 3114 ]
+
       ORACLE_TYPES = {
         :blob=>lambda{|b| Sequel::SQL::Blob.new(b.read)},
         :clob=>lambda(&:read)
       }
 
+      NULL = LiteralString.new('NULL').freeze
+      NULL_RETURNING = [NULL].freeze
+      NULL_RETURNING_BINDING = [[NULL, String].freeze].freeze
+
       # Hash of conversion procs for this database.
       attr_reader :conversion_procs
 
       def connect(server)
         opts = server_opts(server)
-        if opts[:database]
-          dbname = opts[:host] ? \
-            "//#{opts[:host]}#{":#{opts[:port]}" if opts[:port]}/#{opts[:database]}" : opts[:database]
+        if opts[:database] && opts[:host]
+          port = opts[:port] ? ":#{opts[:port]}" : ""
+          dbname = "//#{opts[:host]}#{port}/#{opts[:database]}"
         else
-          dbname = opts[:host]
+          dbname = opts[:database] || opts[:host]
         end
         conn = OCI8.new(opts[:user], opts[:password], dbname, opts[:privilege])
         if prefetch_rows = opts.fetch(:prefetch_rows, 100)
@@ -76,6 +80,44 @@ def freeze
         super
       end
 
+      # Disables automatic use of INSERT ... RETURNING.  You can still use
+      # returning manually to force the use of RETURNING when inserting.
+      #
+      # This is designed for cases where INSERT RETURNING cannot be used,
+      # such as performing DML operations on views with INSTEAD OF triggers
+      #
+      # Note that when this method is used, insert will not return the
+      # primary key of the inserted row, you will have to get the primary
+      # key of the inserted row before inserting via nextval, or after
+      # inserting via currval or lastval (making sure to use the same
+      # database connection for currval or lastval).
+      def disable_insert_returning
+        clone(:disable_insert_returning=>true)
+      end
+
+      # Return primary key for the given table.
+      def primary_key(table)
+        quoted_table = quote_schema_table(table)
+        Sequel.synchronize{return @primary_keys[quoted_table] if @primary_keys.key?(quoted_table)}
+        value, _ = schema(table).find { |_, c| c[:primary_key] }
+        Sequel.synchronize{@primary_keys[quoted_table] = value}
+      end
+
+      RETURNING_TYPES = {:string=>String, :integer=>Integer}.freeze
+      def returning_values(table, columns)
+        quoted_table = quote_schema_table(table)
+        Sequel.synchronize{return @returning_values[quoted_table][columns] if @returning_values[quoted_table].key?(columns)}
+        if columns == NULL_RETURNING
+          values = NULL_RETURNING_BINDING
+        else
+          col_names = columns.map(&:value)
+          values = schema(table).map do |(name, metadata)|
+            [name, RETURNING_TYPES[metadata[:type]]] if col_names.include?(name)
+          end.compact
+        end
+        Sequel.synchronize{@returning_values[quoted_table][columns] = values}
+      end
+
       private
 
       def _execute(type, sql, opts=OPTS, &block)
@@ -87,12 +129,17 @@ def _execute(type, sql, opts=OPTS, &block)
               args = cursor_bind_params(conn, r, args)
               nr = log_connection_yield(sql, conn, args){r.exec}
               r = nr unless block_given?
+            elsif opts[:returning]
+              args = opts[:returning].map {|(_, type)| [nil, type]}
+              r = conn.parse(sql)
+              args = cursor_bind_params(conn, r, args)
+              nr = log_connection_yield(sql, conn, args){r.exec}
             else
               r = log_connection_yield(sql, conn){conn.exec(sql)}
             end
             if block_given?
               yield(r)
-            elsif type == :insert
+            elsif type == :insert && !opts[:returning]
               last_insert_id(conn, opts)
             else
               r
@@ -110,6 +157,10 @@ def adapter_initialize
         @autosequence = @opts[:autosequence]
         @primary_key_sequences = {}
         @conversion_procs = ORACLE_TYPES.dup
+        @primary_keys = {}
+        @returning_values = Hash.new {|h, k| h[k] = {}}
+
+        super
       end
 
       PS_TYPES = {'string'.freeze=>String, 'integer'.freeze=>Integer, 'float'.freeze=>Float,
@@ -131,6 +182,8 @@ def cursor_bind_params(conn, cursor, args)
           end
           if t = PS_TYPES[type]
             cursor.bind_param(i, arg, t)
+          elsif type
+            cursor.bind_param(i, arg, type)
           else
             cursor.bind_param(i, arg, arg.class)
           end
@@ -261,8 +314,10 @@ def schema_parse_table(table, opts=OPTS)
         schema ||= opts[:schema]
         schema_and_table = if ds = opts[:dataset]
           ds.literal(schema ? SQL::QualifiedIdentifier.new(schema, table) : SQL::Identifier.new(table))
+        elsif schema
+          "#{quote_identifier(schema)}.#{quote_identifier(table)}"
         else
-          "#{"#{quote_identifier(schema)}." if schema}#{quote_identifier(table)}"
+          quote_identifier(table)
         end
         table_schema = []
         m = output_identifier_meth(ds)
@@ -327,7 +382,10 @@ class Dataset < Sequel::Dataset
       Database::DatasetClass = self
 
       PREPARED_ARG_PLACEHOLDER = ':'.freeze
-
+      NULL = Database::NULL
+      DUMMY_RETURNING = Sequel.lit(' RETURNING NULL INTO :dummy').freeze
+      def_sql_method(self, :insert, %w'with insert into columns values returning')
+
       # Oracle already supports named bind arguments, so use directly.
       module ArgumentMapper
         include Sequel::Dataset::ArgumentMapper
@@ -361,16 +419,20 @@ def prepared_arg?(k)
       PreparedStatementMethods = prepared_statements_module(:prepare, BindArgumentMethods)
 
       def fetch_rows(sql)
-        execute(sql) do |cursor|
-          cps = db.conversion_procs
-          cols = columns = cursor.get_col_names.map{|c| output_identifier(c)}
-          metadata = cursor.column_metadata
-          cm = cols.zip(metadata).map{|c, m| [c, cps[m.data_type]]}
-          self.columns = columns
-          while r = cursor.fetch
-            row = {}
-            r.zip(cm).each{|v, (c, cp)| row[c] = ((v && cp) ? cp.call(v) : v)}
-            yield row
+        execute(sql, opts) do |cursor|
+          if opts[:returning]
+            yield Hash[*opts[:returning].flat_map.with_index {|(name, _), idx| [name, cursor[idx+1]]}]
+          else
+            cps = db.conversion_procs
+            cols = columns = cursor.get_col_names.map{|c| output_identifier(c)}
+            metadata = cursor.column_metadata
+            cm = cols.zip(metadata).map{|c, m| [c, cps[m.data_type]]}
+            self.columns = columns
+            while r = cursor.fetch
+              row = {}
+              r.zip(cm).each{|v, (c, cp)| row[c] = ((v && cp) ? cp.call(v) : v)}
+              yield row
+            end
           end
         end
         self
@@ -383,6 +445,41 @@ def requires_placeholder_type_specifiers?
         true
       end
 
+      # Oracle supports for all statements.
+      def supports_returning?(type)
+        true
+      end
+
+      # Insert given values into the database.
+      def insert(*values)
+        if @opts[:returning]
+          # Already know which columns to return, let the standard code handle it
+          super
+        elsif @opts[:sql] || @opts[:disable_insert_returning]
+          # Raw SQL used or RETURNING disabled, just use the default behavior
+          super
+        else
+          # Force the use of RETURNING with the primary key value,
+          # unless it has been disabled.
+          returning(insert_pk).insert(*values){|r| return r.values.first}
+        end
+      end
+
+      def insert_returning_sql(sql)
+        if opts[:returning]
+          if opts[:returning][0][0] == NULL
+            sql << DUMMY_RETURNING
+          else
+            sql << Dataset::RETURNING
+            column_list_append(sql, opts[:returning].map(&:first))
+            sql << Dataset::INTO
+            column_list_append(sql, opts[:returning].map {|(c, _)| Sequel.lit(":#{c}") })
+          end
+        end
+      end
+      alias delete_returning_sql insert_returning_sql
+      alias update_returning_sql insert_returning_sql
+
       private
 
       def literal_other_append(sql, v)
@@ -397,6 +494,11 @@ def literal_other_append(sql, v)
         end
       end
 
+      def returning(*values)
+        raise Error, "RETURNING is not supported on #{db.database_type}" unless supports_returning?(:insert)
+        clone(:returning=>db.returning_values(opts[:from].first, values).freeze)
+      end
+
       def prepared_arg_placeholder
         PREPARED_ARG_PLACEHOLDER
       end
@@ -408,6 +510,12 @@ def bound_variable_modules
       def prepared_statement_modules
         [PreparedStatementMethods]
       end
+
+      # Return the primary key to use for RETURNING in an INSERT statement.
+      def insert_pk
+        pk = db.primary_key(opts[:from].first)
+        pk ? Sequel::SQL::Identifier.new(pk) : NULL
+      end
     end
   end
 end

lib/sequel/adapters/shared/firebird.rb

@@ -214,7 +214,7 @@ def empty_from_sql
         DEFAULT_FROM
       end
 
-      def insert_pk(*values)
+      def insert_pk
         pk = db.primary_key(opts[:from].first)
         pk ? Sequel::SQL::Identifier.new(pk) : NULL
       end

lib/sequel/adapters/shared/oracle.rb

@@ -621,6 +621,12 @@ def select_lock_sql(sql)
       def supports_quoted_function_names?
         true
       end
+
+      # Return the primary key to use for RETURNING in an INSERT statement.
+      def insert_pk
+        pk = db.primary_key(opts[:from].first)
+        pk ? Sequel::SQL::Identifier.new(pk) : NULL
+      end
     end
   end
 end

lib/sequel/dataset/actions.rb

@@ -1061,6 +1061,10 @@ def execute(sql, opts=OPTS, &block)
         opts[:server] = @opts[:server] || (@opts[:lock] ? :default : :read_only)
         opts
       end
+      if @opts.key?(:returning)
+        opts = Hash[opts]
+        opts[:returning] = @opts[:returning]
+      end
       db.execute(sql, opts, &block)
     end
 

spec/adapters/oracle_spec.rb

@@ -2,6 +2,15 @@
 
 require File.join(File.dirname(File.expand_path(__FILE__)), 'spec_helper.rb')
 
+def DB.sqls
+  (@sqls ||= [])
+end
+logger = Object.new
+def logger.method_missing(m, msg)
+  DB.sqls << msg
+end
+DB.loggers << logger
+
 describe "An Oracle database" do
   before(:all) do
     DB.create_table!(:items) do
@@ -12,7 +21,7 @@
     end
 
     DB.create_table!(:books) do
-      Integer :id
+      primary_key :id
       String :title, :size => 50
       Integer :category_id
     end
@@ -29,6 +38,10 @@
     end
     @d = DB[:items]
   end
+  before do
+    @db = DB
+    DB.sqls.clear
+  end
   after do
     @d.delete
   end
@@ -86,7 +99,7 @@
   end
 
   it "should provide schema information" do
-    books_schema = [[:id, [:integer, false, true, nil]],
+    books_schema = [[:id, [:integer, true, false, nil]],
       [:title, [:string, false, true, nil]],
       [:category_id, [:integer, false, true, nil]]]
     categories_schema = [[:id, [:integer, false, true, nil]],
@@ -278,14 +291,14 @@
       {:id => 1, :title => 'aaa', :cat_name => 'ruby'},
       {:id => 2, :title => 'bbb', :cat_name => 'ruby'},
       {:id => 3, :title => 'ccc', :cat_name => 'rails'},
-      {:id => 4, :title => 'ddd', :cat_name => nil} 
+      {:id => 4, :title => 'ddd', :cat_name => nil}
     ]
-    
-    @d1.left_outer_join(:categories, :id => :category_id).select(Sequel[:books][:id], :title, :cat_name).reverse_order(Sequel[:books][:id]).limit(2, 0).to_a.must_equal [      
-      {:id => 4, :title => 'ddd', :cat_name => nil}, 
+
+    @d1.left_outer_join(:categories, :id => :category_id).select(Sequel[:books][:id], :title, :cat_name).reverse_order(Sequel[:books][:id]).limit(2, 0).to_a.must_equal [
+      {:id => 4, :title => 'ddd', :cat_name => nil},
       {:id => 3, :title => 'ccc', :cat_name => 'rails'}
-    ]      
-  end  
+    ]
+  end
 
   it "should allow columns to be renamed" do
     @d1 = DB[:books]
@@ -310,4 +323,18 @@
   it "#lock_style should accept symbols" do
     DB[:books].lock_style(:update).sql.must_equal 'SELECT * FROM "BOOKS" FOR UPDATE'
   end
+
+  unless defined? JRUBY_VERSION
+    it "should use INSERT RETURNING statement by default" do
+      @db[:books].delete
+
+      log do
+        @db[:books].insert(:title => 'aaa').must_equal 1
+      end
+
+      check_sqls do
+        @db.sqls.last.must_equal "INSERT INTO \"BOOKS\" (\"TITLE\") VALUES ('aaa') RETURNING \"ID\" INTO :id; [nil]"
+      end
+    end
+  end
 end