Add schema argument to AutomapBase.prepare()

This allows automap to reflect tables from a schema other than the
default without the need to resort to calling MetaData.reflect
directly.

Change-Id: Ie73cb113bd6d115555c09c5efc33d27ad2c9c512
Pull-request: https://github.com/zzzeek/sqlalchemy/pull/237

3 files changed, 55 insertions, 1 deletions

diff --git a/doc/build/changelog/changelog_11.rst b/doc/build/changelog/changelog_11.rst
index 898b8a0ba..8e7c962ce 100644
--- a/doc/build/changelog/changelog_11.rst
+++ b/doc/build/changelog/changelog_11.rst
@@ -31,6 +31,15 @@
         request courtesy Benjamin Bertrand.
 
     .. change::
+        :tags: orm, feature
+        :pullreq: github:237
+
+        Added :paramref:`.AutomapBase.prepare.schema` to the
+        :meth:`.AutomapBase.prepare` method, to indicate which schema
+        tables should be reflected from if not the default schema.
+        Pull request courtesy Josh Marlow.
+
+    .. change::
         :tags: bug, mssql
         :pullreq: bitbucket:58
 
diff --git a/lib/sqlalchemy/ext/automap.py b/lib/sqlalchemy/ext/automap.py
index bb1c5da3d..23ad56b76 100644
--- a/lib/sqlalchemy/ext/automap.py
+++ b/lib/sqlalchemy/ext/automap.py
@@ -695,6 +695,7 @@ class AutomapBase(object):
             cls,
             engine=None,
             reflect=False,
+            schema=None,
             classname_for_table=classname_for_table,
             collection_class=list,
             name_for_scalar_relationship=name_for_scalar_relationship,
@@ -735,10 +736,19 @@ class AutomapBase(object):
          when a new :func:`.relationship` object is created that represents a
          collection.  Defaults to ``list``.
 
+        :param schema: When present in conjunction with the
+         :paramref:`.AutomapBase.prepare.reflect` flag, is passed to
+         :meth:`.MetaData.reflect` to indicate the primary schema where tables
+         should be reflected from.  When omitted, the default schema in use
+         by the database connection is used.
+
+         .. versionadded:: 1.1
+
         """
         if reflect:
             cls.metadata.reflect(
                 engine,
+                schema=schema,
                 extend_existing=True,
                 autoload_replace=False
             )
diff --git a/test/ext/test_automap.py b/test/ext/test_automap.py
index 0a57b9caa..0c5c4e9dd 100644
--- a/test/ext/test_automap.py
+++ b/test/ext/test_automap.py
@@ -3,7 +3,7 @@ from ..orm._fixtures import FixtureTest
 from sqlalchemy.ext.automap import automap_base
 from sqlalchemy.orm import relationship, interfaces, configure_mappers
 from sqlalchemy.ext.automap import generate_relationship
-from sqlalchemy.testing.mock import Mock
+from sqlalchemy.testing.mock import Mock, patch
 from sqlalchemy import String, Integer, ForeignKey
 from sqlalchemy import testing
 from sqlalchemy.testing.schema import Table, Column
@@ -71,6 +71,41 @@ class AutomapTest(fixtures.MappedTest):
         n1.nodes_collection.append(n2)
         assert n2.nodes is n1
 
+    def test_prepare_accepts_optional_schema_arg(self):
+        """
+        The underlying reflect call accepts an optional schema argument.
+        This is for determining which database schema to load.
+        This test verifies that prepare can accept an optiona schema argument
+        and pass it to reflect.
+        """
+        Base = automap_base(metadata=self.metadata)
+        engine_mock = Mock()
+        with patch.object(Base.metadata, "reflect") as reflect_mock:
+            Base.prepare(engine_mock, reflect=True, schema="some_schema")
+            reflect_mock.assert_called_once_with(
+                engine_mock,
+                schema="some_schema",
+                extend_existing=True,
+                autoload_replace=False,
+            )
+
+    def test_prepare_defaults_to_no_schema(self):
+        """
+        The underlying reflect call accepts an optional schema argument.
+        This is for determining which database schema to load.
+        This test verifies that prepare passes a default None if no schema is provided.
+        """
+        Base = automap_base(metadata=self.metadata)
+        engine_mock = Mock()
+        with patch.object(Base.metadata, "reflect") as reflect_mock:
+            Base.prepare(engine_mock, reflect=True)
+            reflect_mock.assert_called_once_with(
+                engine_mock,
+                schema=None,
+                extend_existing=True,
+                autoload_replace=False,
+            )
+
     def test_naming_schemes(self):
         Base = automap_base(metadata=self.metadata)