[ticket:1893] implementation

4 files changed, 115 insertions, 12 deletions

diff --git a/lib/sqlalchemy/ext/declarative.py b/lib/sqlalchemy/ext/declarative.py
index dd2df63d3..53438dc55 100755
--- a/lib/sqlalchemy/ext/declarative.py
+++ b/lib/sqlalchemy/ext/declarative.py
@@ -924,7 +924,7 @@ Mapped instances then make usage of
 
 """
 
-from sqlalchemy.schema import Table, Column, MetaData
+from sqlalchemy.schema import Table, Column, MetaData, _get_table_key
 from sqlalchemy.orm import synonym as _orm_synonym, mapper,\
                                 comparable_property, class_mapper
 from sqlalchemy.orm.interfaces import MapperProperty
@@ -1257,8 +1257,8 @@ class DeclarativeMeta(type):
 class _GetColumns(object):
     def __init__(self, cls):
         self.cls = cls
-    def __getattr__(self, key):
         
+    def __getattr__(self, key):
         mapper = class_mapper(self.cls, compile=False)
         if mapper:
             prop = mapper.get_property(key, raiseerr=False)
@@ -1273,7 +1273,16 @@ class _GetColumns(object):
                             " directly to a Column)." % key)
         return getattr(self.cls, key)
 
-
+class _GetTable(object):
+    def __init__(self, key, metadata):
+        self.key = key
+        self.metadata = metadata
+    
+    def __getattr__(self, key):
+        return self.metadata.tables[
+                _get_table_key(key, self.key)
+            ]
+        
 def _deferred_relationship(cls, prop):
     def resolve_arg(arg):
         import sqlalchemy
@@ -1283,6 +1292,8 @@ def _deferred_relationship(cls, prop):
                 return _GetColumns(cls._decl_class_registry[key])
             elif key in cls.metadata.tables:
                 return cls.metadata.tables[key]
+            elif key in cls.metadata._schemas:
+                return _GetTable(key, cls.metadata)
             else:
                 return sqlalchemy.__dict__[key]
 
diff --git a/lib/sqlalchemy/schema.py b/lib/sqlalchemy/schema.py
index a332cec36..cb0cde4df 100644
--- a/lib/sqlalchemy/schema.py
+++ b/lib/sqlalchemy/schema.py
@@ -206,12 +206,13 @@ class Table(SchemaItem, expression.TableClause):
             if mustexist:
                 raise exc.InvalidRequestError(
                     "Table '%s' not defined" % (key))
-            metadata.tables[key] = table = object.__new__(cls)
+            table = object.__new__(cls)                    
+            metadata._add_table(name, schema, table)
             try:
                 table._init(name, metadata, *args, **kw)
                 return table
             except:
-                metadata.tables.pop(key)
+                metadata._remove_table(name, schema)
                 raise
                 
     def __init__(self, *args, **kw):
@@ -406,7 +407,7 @@ class Table(SchemaItem, expression.TableClause):
         self.ddl_listeners[event].append(listener)
 
     def _set_parent(self, metadata):
-        metadata.tables[_get_table_key(self.name, self.schema)] = self
+        metadata._add_table(self.name, self.schema, self)
         self.metadata = metadata
 
     def get_children(self, column_collections=True, 
@@ -1938,7 +1939,8 @@ class MetaData(SchemaItem):
           ``MetaData``.
 
         """
-        self.tables = {}
+        self.tables = util.frozendict()
+        self._schemas = set()
         self.bind = bind
         self.metadata = self
         self.ddl_listeners = util.defaultdict(list)
@@ -1957,6 +1959,20 @@ class MetaData(SchemaItem):
             table_or_key = table_or_key.key
         return table_or_key in self.tables
 
+    def _add_table(self, name, schema, table):
+        key = _get_table_key(name, schema)
+        dict.__setitem__(self.tables, key, table)
+        if schema:
+            self._schemas.add(schema)
+    
+    def _remove_table(self, name, schema):
+        key = _get_table_key(name, schema)
+        dict.pop(self.tables, key, None)
+        if self._schemas:
+            self._schemas = set([t.schema 
+                                for t in self.tables.values() 
+                                if t.schema is not None])
+        
     def __getstate__(self):
         return {'tables': self.tables}
 
@@ -1991,15 +2007,14 @@ class MetaData(SchemaItem):
 
     def clear(self):
         """Clear all Table objects from this MetaData."""
-        # TODO: why have clear()/remove() but not all
-        # other accesors/mutators for the tables dict ?
-        self.tables.clear()
 
+        dict.clear(self.tables)
+        self._schemas.clear()
+        
     def remove(self, table):
         """Remove the given Table object from this MetaData."""
         
-        # TODO: scan all other tables and remove FK _column
-        del self.tables[table.key]
+        self._remove_table(table.name, table.schema)
 
     @property
     def sorted_tables(self):
diff --git a/test/engine/test_metadata.py b/test/engine/test_metadata.py
index b2250c808..59aa4c354 100644
--- a/test/engine/test_metadata.py
+++ b/test/engine/test_metadata.py
@@ -77,6 +77,51 @@ class MetaDataTest(TestBase, ComparesTables):
             t = Table('foo%d' % i, m, cx)
         eq_(msgs, ['attach foo0.foo', 'attach foo1.foo', 'attach foo2.foo'])
         
+    def test_schema_collection_add(self):
+        metadata = MetaData()
+        
+        t1 = Table('t1', metadata, Column('x', Integer), schema='foo')
+        t2 = Table('t2', metadata, Column('x', Integer), schema='bar')
+        t3 = Table('t3', metadata, Column('x', Integer))
+        
+        eq_(metadata._schemas, set(['foo', 'bar']))
+        eq_(len(metadata.tables), 3)
+    
+    def test_schema_collection_remove(self):
+        metadata = MetaData()
+        
+        t1 = Table('t1', metadata, Column('x', Integer), schema='foo')
+        t2 = Table('t2', metadata, Column('x', Integer), schema='bar')
+        t3 = Table('t3', metadata, Column('x', Integer), schema='bar')
+        
+        metadata.remove(t3)
+        eq_(metadata._schemas, set(['foo', 'bar']))
+        eq_(len(metadata.tables), 2)
+
+        metadata.remove(t1)
+        eq_(metadata._schemas, set(['bar']))
+        eq_(len(metadata.tables), 1)
+    
+    def test_schema_collection_remove_all(self):
+        metadata = MetaData()
+        
+        t1 = Table('t1', metadata, Column('x', Integer), schema='foo')
+        t2 = Table('t2', metadata, Column('x', Integer), schema='bar')
+
+        metadata.clear()
+        eq_(metadata._schemas, set())
+        eq_(len(metadata.tables), 0)
+    
+    def test_metadata_tables_immutable(self):
+        metadata = MetaData()
+        
+        t1 = Table('t1', metadata, Column('x', Integer))
+        assert 't1' in metadata.tables
+        
+        assert_raises(
+            AttributeError,
+            lambda: metadata.tables.pop('t1')
+        )
         
     def test_dupe_tables(self):
         metadata = MetaData()
diff --git a/test/ext/test_declarative.py b/test/ext/test_declarative.py
index 72e2edf30..64dc81496 100644
--- a/test/ext/test_declarative.py
+++ b/test/ext/test_declarative.py
@@ -299,6 +299,38 @@ class DeclarativeTest(DeclarativeTestBase):
         assert class_mapper(User).get_property('props').secondary \
             is user_to_prop
 
+    def test_string_dependency_resolution_schemas(self):
+        Base = decl.declarative_base()
+        
+        class User(Base):
+
+            __tablename__ = 'users'
+            __table_args__ = {'schema':'fooschema'}
+            
+            id = Column(Integer, primary_key=True)
+            name = Column(String(50))
+            props = relationship('Prop', secondary='fooschema.user_to_prop',
+                         primaryjoin='User.id==fooschema.user_to_prop.c.user_id',
+                         secondaryjoin='fooschema.user_to_prop.c.prop_id==Prop.id', 
+                         backref='users')
+
+        class Prop(Base):
+
+            __tablename__ = 'props'
+            __table_args__ = {'schema':'fooschema'}
+
+            id = Column(Integer, primary_key=True)
+            name = Column(String(50))
+
+        user_to_prop = Table('user_to_prop', Base.metadata,
+                     Column('user_id', Integer, ForeignKey('fooschema.users.id')), 
+                     Column('prop_id',Integer, ForeignKey('fooschema.props.id')),
+                     schema='fooschema')
+        compile_mappers()
+        
+        assert class_mapper(User).get_property('props').secondary \
+            is user_to_prop
+
     def test_uncompiled_attributes_in_relationship(self):
 
         class Address(Base, ComparableEntity):