Fix handling of schema set by search-path in pg.

Author: coleifer

CHANGELOG.md

@@ -29,6 +29,10 @@ Backwards-incompatible:
 
 Improvements:
 
+* Postgres introspection (`get_tables()`, `get_columns()`, `get_indexes()`,
+  `get_primary_keys()`, `get_foreign_keys()`, `get_views()`) with no `schema`
+  now follows the search path via `current_schema()` instead of assuming
+  `public`, matching MySQL's `DATABASE()` behavior.
 * Add `SchemaMigrator(db, schema=...)`, which qualifies every table name w/the
   given schema (or database on MySQL). Not supported for SQLite, as the
   table-rebuild rewrites DDL straight from `sqlite_master`.

peewee.py

@@ -4795,14 +4795,16 @@ def begin(self, isolation_level=None):
 
     def get_tables(self, schema=None):
         query = ('SELECT tablename FROM pg_catalog.pg_tables '
-                 'WHERE schemaname = %s ORDER BY tablename')
-        cursor = self.execute_sql(query, (schema or 'public',))
+                 'WHERE schemaname = COALESCE(%s, current_schema()) '
+                 'ORDER BY tablename')
+        cursor = self.execute_sql(query, (schema,))
         return [table for table, in cursor.fetchall()]
 
     def get_views(self, schema=None):
         query = ('SELECT viewname, definition FROM pg_catalog.pg_views '
-                 'WHERE schemaname = %s ORDER BY viewname')
-        cursor = self.execute_sql(query, (schema or 'public',))
+                 'WHERE schemaname = COALESCE(%s, current_schema()) '
+                 'ORDER BY viewname')
+        cursor = self.execute_sql(query, (schema,))
         return [ViewMetadata(view_name, sql.strip(' \t;'))
                 for (view_name, sql) in cursor.fetchall()]
 
@@ -4822,9 +4824,10 @@ def get_indexes(self, table, schema=None):
                 idxs.tablename = t.relname
                 AND idxs.indexname = i.relname
                 AND idxs.schemaname = n.nspname)
-            WHERE t.relname = %s AND t.relkind = %s AND n.nspname = %s
+            WHERE t.relname = %s AND t.relkind = %s
+                AND n.nspname = COALESCE(%s, current_schema())
             ORDER BY idx.indisunique DESC, i.relname;"""
-        cursor = self.execute_sql(query, (table, 'r', schema or 'public'))
+        cursor = self.execute_sql(query, (table, 'r', schema))
         return [IndexMetadata(name, sql.rstrip(' ;'),
                               [unqesc(c) for c in cols], unique, table)
                 for name, sql, unique, cols in cursor.fetchall()]
@@ -4842,10 +4845,11 @@ def get_columns(self, table, schema=None):
                 ON (t.relname = c.table_name AND t.relnamespace = n.oid)
             INNER JOIN pg_catalog.pg_attribute AS a
                 ON (a.attrelid = t.oid AND a.attname = c.column_name)
-            WHERE c.table_name = %s AND c.table_schema = %s
+            WHERE c.table_name = %s
+                AND c.table_schema = COALESCE(%s, current_schema())
                 AND NOT a.attisdropped
             ORDER BY c.ordinal_position"""
-        cursor = self.execute_sql(query, (table, schema or 'public'))
+        cursor = self.execute_sql(query, (table, schema))
         pks = set(self.get_primary_keys(table, schema))
         def is_ident(ident, df):
             return ident == 'YES' or (df or '').startswith(
@@ -4865,9 +4869,9 @@ def get_primary_keys(self, table, schema=None):
             WHERE
                 tc.constraint_type = %s AND
                 tc.table_name = %s AND
-                tc.table_schema = %s"""
+                tc.table_schema = COALESCE(%s, current_schema())"""
         ctype = 'PRIMARY KEY'
-        cursor = self.execute_sql(query, (ctype, table, schema or 'public'))
+        cursor = self.execute_sql(query, (ctype, table, schema))
         return [pk for pk, in cursor.fetchall()]
 
     def get_foreign_keys(self, table, schema=None):
@@ -4890,8 +4894,8 @@ def get_foreign_keys(self, table, schema=None):
             WHERE
                 tc.constraint_type = 'FOREIGN KEY' AND
                 tc.table_name = %s AND
-                tc.table_schema = %s"""
-        cursor = self.execute_sql(sql, (table, schema or 'public'))
+                tc.table_schema = COALESCE(%s, current_schema())"""
+        cursor = self.execute_sql(sql, (table, schema))
         return [ForeignKeyMetadata(row[0], row[1], row[2], table, row[3],
                                    row[4], row[5])
                 for row in cursor.fetchall()]

tests/db_tests.py

@@ -756,6 +756,46 @@ def test_quoted_name(self):
         self.assertEqual(idx.columns, ['da"ta'])
 
 
+@requires_postgresql
+class TestIntrospectionSearchPath(DatabaseTestCase):
+    """With no schema given, introspection follows the search path."""
+    schema = 'ispath'
+
+    def setUp(self):
+        super(TestIntrospectionSearchPath, self).setUp()
+        self.execute('DROP SCHEMA IF EXISTS %s CASCADE' % self.schema)
+        self.execute('CREATE SCHEMA %s' % self.schema)
+        self.execute('CREATE TABLE %s.parent (id SERIAL PRIMARY KEY)'
+                     % self.schema)
+        self.execute('CREATE TABLE %s.child (id SERIAL PRIMARY KEY, '
+                     'parent_id INTEGER REFERENCES %s.parent (id), '
+                     'name TEXT)' % (self.schema, self.schema))
+        self.execute('CREATE INDEX child_name ON %s.child (name)'
+                     % self.schema)
+        self.execute('CREATE VIEW %s.child_names AS SELECT name FROM '
+                     '%s.child' % (self.schema, self.schema))
+        self.execute('SET search_path TO %s' % self.schema)
+
+    def tearDown(self):
+        try:
+            self.execute('DROP SCHEMA IF EXISTS %s CASCADE' % self.schema)
+        finally:
+            super(TestIntrospectionSearchPath, self).tearDown()
+
+    def test_search_path(self):
+        db = self.database
+        self.assertEqual(db.get_tables(), ['child', 'parent'])
+        self.assertEqual([v.name for v in db.get_views()], ['child_names'])
+        self.assertEqual([c.name for c in db.get_columns('child')],
+                         ['id', 'parent_id', 'name'])
+        self.assertEqual(db.get_primary_keys('child'), ['id'])
+        self.assertEqual([(fk.column, fk.dest_table)
+                          for fk in db.get_foreign_keys('child')],
+                         [('parent_id', 'parent')])
+        self.assertTrue('child_name' in
+                        [i.name for i in db.get_indexes('child')])
+
+
 # ===========================================================================
 # Thread safety
 # ===========================================================================