The AutomapBase's prepare method now accepts an optional schema argument to pass to reflectpr/237

2 files changed, 38 insertions, 1 deletions

diff --git a/lib/sqlalchemy/ext/automap.py b/lib/sqlalchemy/ext/automap.py
index 023d11ca8..8b3147010 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,
@@ -739,6 +740,7 @@ class AutomapBase(object):
         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)