Improve PostgreSQL adapter schema-awareness

* table_exists? scoped by schema search path unless schema is explicitly named. Added tests and doc to clarify the behaviour
* extract_schema_and_table tests and implementation extended to cover all cases
* primary_key does not ignore schema information
* add current_schema and schema_exists? methods
* more robust table referencing in insert_sql and sql_for_insert methods

3 files changed, 173 insertions, 27 deletions

diff --git a/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb b/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb
index 3e390ba994..2a3ee33e3e 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,33 +666,48 @@ module ActiveRecord
         SQL
       end
 
+      # Returns true of 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, schema] if schema
+        binds << [nil, schema.gsub(/(^"|"$)/,'')] 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_tables
+          WHERE tablename = $1
+          AND schemaname = #{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
+      # 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
 
-        if name =~ /^"/ # Handle quoted table names
-          table  = name
-          schema = nil
-        end
+      # Returns an array of [schema_name, table_name] extracted from +name+.
+      # The schema_name will be nil if not provided in +name+.
+      # Quotes are preserved in the schema and table name components if provided.
+      # Valid combinations for quoting the schema and table names:
+      #
+      # - table_name
+      # - "table.name"
+      # - schema_name.table_name
+      # - schema_name."table.name"
+      # - "schema.name".table_name
+      # - "schema.name"."table_name"
+      # - "schema.name"."table.name"
+      def extract_schema_and_table(name)
+        name[/([^"\.\s]+|"[^"]+")(?:\.([^"\.\s]+|"[^"]*"))?/]
+        table, schema = [$1,$2].compact.reverse
         [schema, table]
       end
 
@@ -742,6 +758,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]
@@ -843,7 +864,7 @@ module ActiveRecord
 
       # Returns just a table's primary key
       def primary_key(table)
-        row = exec_query(<<-end_sql, 'SCHEMA', [[nil, table]]).rows.first
+        row = exec_query(<<-end_sql, 'SCHEMA', [[nil, quote_table_name(table)]]).rows.first
           SELECT DISTINCT(attr.attname)
           FROM pg_attribute attr
           INNER JOIN pg_depend dep ON attr.attrelid = dep.refobjid AND attr.attnum = dep.refobjsubid
@@ -1080,6 +1101,11 @@ module ActiveRecord
           end
         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
diff --git a/activerecord/test/cases/adapters/postgresql/postgresql_adapter_test.rb b/activerecord/test/cases/adapters/postgresql/postgresql_adapter_test.rb
index 7c49236854..b113267dca 100644
--- a/activerecord/test/cases/adapters/postgresql/postgresql_adapter_test.rb
+++ b/activerecord/test/cases/adapters/postgresql/postgresql_adapter_test.rb
@@ -10,6 +10,45 @@ module ActiveRecord
         @connection.exec_query('create table ex(id serial primary key, number integer, data character varying(255))')
       end
 
+      def test_primary_key
+        assert_equal 'id',@connection.primary_key('ex')
+      end
+
+      def test_non_standard_primary_key
+        @connection.exec_query('drop table if exists ex')
+        @connection.exec_query('create table ex(data character varying(255) primary key)')
+        assert_equal 'data', @connection.primary_key('ex')
+      end
+
+      def test_primary_key_returns_nil_for_no_pk
+        @connection.exec_query('drop table if exists ex')
+        @connection.exec_query('create table ex(id integer)')
+        assert_nil @connection.primary_key('ex')
+      end
+
+      def test_primary_key_raises_error_if_table_not_found
+        assert_raises(ActiveRecord::StatementInvalid) do
+          @connection.primary_key('unobtainium')
+        end
+      end
+
+      def test_insert_sql_with_proprietary_returning_clause
+        id = @connection.insert_sql("insert into ex (number) values(5150)", nil, "number")
+        assert_equal "5150", id
+      end
+
+      def test_insert_sql_with_quoted_schema_and_table_name
+        id = @connection.insert_sql('insert into "public"."ex" (number) values(5150)')
+        expect = @connection.query('select max(id) from ex').first.first
+        assert_equal expect, id
+      end
+
+      def test_insert_sql_with_no_space_after_table_name
+        id = @connection.insert_sql("insert into ex(number) values(5150)")
+        expect = @connection.query('select max(id) from ex').first.first
+        assert_equal expect, id
+      end
+
       def test_serial_sequence
         assert_equal 'public.accounts_id_seq',
           @connection.serial_sequence('accounts', 'id')
diff --git a/activerecord/test/cases/adapters/postgresql/schema_test.rb b/activerecord/test/cases/adapters/postgresql/schema_test.rb
index a5c3e69af9..d4797e1680 100644
--- a/activerecord/test/cases/adapters/postgresql/schema_test.rb
+++ b/activerecord/test/cases/adapters/postgresql/schema_test.rb
@@ -20,6 +20,7 @@ class SchemaTest < ActiveRecord::TestCase
     'email character varying(50)',
     'moment timestamp without time zone default now()'
   ]
+  PK_TABLE_NAME = 'table_with_pk'
 
   class Thing1 < ActiveRecord::Base
     set_table_name "test_schema.things"
@@ -37,6 +38,10 @@ class SchemaTest < ActiveRecord::TestCase
     set_table_name 'test_schema."Things"'
   end
 
+  class PrimaryKeyTestHarness < ActiveRecord::Base
+    set_table_name 'test_schema.pktest'
+  end
+
   def setup
     @connection = ActiveRecord::Base.connection
     @connection.execute "CREATE SCHEMA #{SCHEMA_NAME} CREATE TABLE #{TABLE_NAME} (#{COLUMNS.join(',')})"
@@ -49,6 +54,7 @@ class SchemaTest < ActiveRecord::TestCase
     @connection.execute "CREATE INDEX #{INDEX_B_NAME} ON #{SCHEMA2_NAME}.#{TABLE_NAME}  USING btree (#{INDEX_B_COLUMN_S2});"
     @connection.execute "CREATE INDEX #{INDEX_C_NAME} ON #{SCHEMA_NAME}.#{TABLE_NAME}  USING gin (#{INDEX_C_COLUMN});"
     @connection.execute "CREATE INDEX #{INDEX_C_NAME} ON #{SCHEMA2_NAME}.#{TABLE_NAME}  USING gin (#{INDEX_C_COLUMN});"
+    @connection.execute "CREATE TABLE #{SCHEMA_NAME}.#{PK_TABLE_NAME} (id serial primary key)"
   end
 
   def teardown
@@ -63,12 +69,30 @@ class SchemaTest < ActiveRecord::TestCase
     end
   end
 
+  def test_table_exists_when_on_schema_search_path
+    with_schema_search_path(SCHEMA_NAME) do
+      assert(@connection.table_exists?(TABLE_NAME), "table should exist and be found")
+    end
+  end
+
+  def test_table_exists_when_not_on_schema_search_path
+    with_schema_search_path('PUBLIC') do
+      assert(!@connection.table_exists?(TABLE_NAME), "table exists but should not be found")
+    end
+  end
+
   def test_table_exists_wrong_schema
     assert(!@connection.table_exists?("foo.things"), "table should not exist")
   end
 
-  def test_table_exists_quoted_table
-    assert(@connection.table_exists?('"things.table"'), "table should exist")
+  def test_table_exists_quoted_names
+    [ %("#{SCHEMA_NAME}"."#{TABLE_NAME}"), %(#{SCHEMA_NAME}."#{TABLE_NAME}"), %(#{SCHEMA_NAME}."#{TABLE_NAME}")].each do |given|
+      assert(@connection.table_exists?(given), "table should exist when specified as #{given}")
+    end
+    with_schema_search_path(SCHEMA_NAME) do
+      given = %("#{TABLE_NAME}")
+      assert(@connection.table_exists?(given), "table should exist when specified as #{given}")
+    end
   end
 
   def test_with_schema_prefixed_table_name
@@ -164,6 +188,63 @@ class SchemaTest < ActiveRecord::TestCase
     ActiveRecord::Base.connection.schema_search_path = "public"
   end
 
+  def test_primary_key_with_schema_specified
+    [ %("#{SCHEMA_NAME}"."#{PK_TABLE_NAME}"), %(#{SCHEMA_NAME}."#{PK_TABLE_NAME}"), %(#{SCHEMA_NAME}."#{PK_TABLE_NAME}")].each do |given|
+      assert_equal 'id', @connection.primary_key(given), "primary key should be found when table referenced as #{given}"
+    end
+  end
+
+  def test_primary_key_assuming_schema_search_path
+    with_schema_search_path(SCHEMA_NAME) do
+      assert_equal 'id', @connection.primary_key(PK_TABLE_NAME), "primary key should be found"
+    end
+  end
+
+  def test_primary_key_raises_error_if_table_not_found_on_schema_search_path
+    with_schema_search_path(SCHEMA2_NAME) do
+      assert_raises(ActiveRecord::StatementInvalid) do
+        @connection.primary_key(PK_TABLE_NAME)
+      end
+    end
+  end
+
+  def test_extract_schema_and_table
+    {
+      %(table_name)            => [nil,'table_name'],
+      %("table.name")          => [nil,'"table.name"'],
+      %(schema.table_name)     => %w{schema table_name},
+      %("schema".table_name)   => %w{"schema" table_name},
+      %(schema."table_name")   => %w{schema "table_name"},
+      %("schema"."table_name") => %w{"schema" "table_name"},
+      %("even spaces".table)   => ['"even spaces"','table'],
+      %(schema."table.name")   => %w{schema "table.name"}
+    }.each do |given,expect|
+      assert_equal expect, @connection.extract_schema_and_table(given)
+    end
+  end
+
+  def test_current_schema
+    {
+      %('$user',public) => 'public',
+      SCHEMA_NAME => SCHEMA_NAME,
+      %(#{SCHEMA2_NAME},#{SCHEMA_NAME},public) => SCHEMA2_NAME,
+      %(public,#{SCHEMA2_NAME},#{SCHEMA_NAME}) => 'public'
+    }.each do |given,expect|
+      with_schema_search_path(given) { assert_equal expect, @connection.current_schema }
+    end
+  end
+
+  def test_schema_exists?
+    {
+      'public'     => true,
+      SCHEMA_NAME  => true,
+      SCHEMA2_NAME => true,
+      'darkside'   => false
+    }.each do |given,expect|
+      assert_equal expect, @connection.schema_exists?(given)
+    end
+  end
+
   private
     def columns(table_name)
       @connection.send(:column_definitions, table_name).map do |name, type, default|