Merge pull request #1632 from tardate/pg_schema_fu

Improving PostgreSQL adapter schema-awareness

1 files changed, 72 insertions, 47 deletions

diff --git a/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb b/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb
index 1a24ed1375..6d638628df 100644
--- a/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb
+++ b/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb
@@ -466,10 +466,11 @@ module ActiveRecord
 
       # Executes an INSERT query and returns the new record's ID
       def insert_sql(sql, name = nil, pk = nil, id_value = nil, sequence_name = nil)
-        # Extract the table from the insert sql. Yuck.
-        _, table = extract_schema_and_table(sql.split(" ", 4)[2])
-
-        pk ||= primary_key(table)
+        unless pk
+          # Extract the table from the insert sql. Yuck.
+          table_ref = extract_table_ref_from_insert_sql(sql)
+          pk = primary_key(table_ref) if table_ref
+        end
 
         if pk
           select_value("#{sql} RETURNING #{quote_column_name(pk)}")
@@ -565,9 +566,9 @@ module ActiveRecord
 
       def sql_for_insert(sql, pk, id_value, sequence_name, binds)
         unless pk
-          _, table = extract_schema_and_table(sql.split(" ", 4)[2])
-
-          pk = primary_key(table)
+          # Extract the table from the insert sql. Yuck.
+          table_ref = extract_table_ref_from_insert_sql(sql)
+          pk = primary_key(table_ref) if table_ref
         end
 
         sql = "#{sql} RETURNING #{quote_column_name(pk)}" if pk
@@ -665,34 +666,33 @@ module ActiveRecord
         SQL
       end
 
+      # Returns true if table exists.
+      # If the schema is not specified as part of +name+ then it will only find tables within
+      # the current schema search path (regardless of permissions to access tables in other schemas)
       def table_exists?(name)
         schema, table = extract_schema_and_table(name.to_s)
+        return false unless table
 
-        binds = [[nil, table.gsub(/(^"|"$)/,'')]]
+        binds = [[nil, table]]
         binds << [nil, schema] if schema
 
         exec_query(<<-SQL, 'SCHEMA', binds).rows.first[0].to_i > 0
-            SELECT COUNT(*)
-            FROM pg_tables
-            WHERE tablename = $1
-            #{schema ? "AND schemaname = $2" : ''}
+          SELECT COUNT(*)
+          FROM pg_class c
+          LEFT JOIN pg_namespace n ON n.oid = c.relnamespace
+          WHERE c.relkind in ('v','r')
+          AND c.relname = $1
+          AND n.nspname = #{schema ? '$2' : 'ANY (current_schemas(false))'}
         SQL
       end
 
-      # Extracts the table and schema name from +name+
-      def extract_schema_and_table(name)
-        schema, table = name.split('.', 2)
-
-        unless table # A table was provided without a schema
-          table  = schema
-          schema = nil
-        end
-
-        if name =~ /^"/ # Handle quoted table names
-          table  = name
-          schema = nil
-        end
-        [schema, table]
+      # Returns true if schema exists.
+      def schema_exists?(name)
+        exec_query(<<-SQL, 'SCHEMA', [[nil, name]]).rows.first[0].to_i > 0
+          SELECT COUNT(*)
+          FROM pg_namespace
+          WHERE nspname = $1
+        SQL
       end
 
       # Returns an array of indexes for the given table.
@@ -742,6 +742,11 @@ module ActiveRecord
         query('select current_database()')[0][0]
       end
 
+      # Returns the current schema name.
+      def current_schema
+        query('SELECT current_schema', 'SCHEMA')[0][0]
+      end
+
       # Returns the current database encoding format.
       def encoding
         query(<<-end_sql)[0][0]
@@ -960,27 +965,27 @@ module ActiveRecord
         end
 
       private
-      def exec_no_cache(sql, binds)
-        @connection.async_exec(sql)
-      end
-
-      def exec_cache(sql, binds)
-        unless @statements.key? sql
-          nextkey = "a#{@statements.length + 1}"
-          @connection.prepare nextkey, sql
-          @statements[sql] = nextkey
+        def exec_no_cache(sql, binds)
+          @connection.async_exec(sql)
         end
 
-        key = @statements[sql]
+        def exec_cache(sql, binds)
+          unless @statements.key? sql
+            nextkey = "a#{@statements.length + 1}"
+            @connection.prepare nextkey, sql
+            @statements[sql] = nextkey
+          end
 
-        # Clear the queue
-        @connection.get_last_result
-        @connection.send_query_prepared(key, binds.map { |col, val|
-          type_cast(val, col)
-        })
-        @connection.block
-        @connection.get_last_result
-      end
+          key = @statements[sql]
+
+          # Clear the queue
+          @connection.get_last_result
+          @connection.send_query_prepared(key, binds.map { |col, val|
+            type_cast(val, col)
+          })
+          @connection.block
+          @connection.get_last_result
+        end
 
         # The internal PostgreSQL identifier of the money data type.
         MONEY_COLUMN_TYPE_OID = 790 #:nodoc:
@@ -1080,9 +1085,29 @@ module ActiveRecord
           end
         end
 
-      def table_definition
-        TableDefinition.new(self)
-      end
+        # Returns an array of <tt>[schema_name, table_name]</tt> extracted from +name+.
+        # +schema_name+ is nil if not specified in +name+.
+        # +schema_name+ and +table_name+ exclude surrounding quotes (regardless of whether provided in +name+)
+        # +name+ supports the range of schema/table references understood by PostgreSQL, for example:
+        #
+        # * <tt>table_name</tt>
+        # * <tt>"table.name"</tt>
+        # * <tt>schema_name.table_name</tt>
+        # * <tt>schema_name."table.name"</tt>
+        # * <tt>"schema.name"."table name"</tt>
+        def extract_schema_and_table(name)
+          table, schema = name.scan(/[^".\s]+|"[^"]*"/)[0..1].collect{|m| m.gsub(/(^"|"$)/,'') }.reverse
+          [schema, table]
+        end
+
+        def extract_table_ref_from_insert_sql(sql)
+          sql[/into\s+([^\(]*).*values\s*\(/i]
+          $1.strip if $1
+        end
+
+        def table_definition
+          TableDefinition.new(self)
+        end
     end
   end
 end