From ae363b0663a8b67f39d2970a6fe0626c54c1478c Mon Sep 17 00:00:00 2001
From: Mike Bayer <mike_mp@zzzcomputing.com>
Date: Sun, 17 Apr 2011 13:53:21 -0400
Subject: [PATCH] - Before/after attach events for PrimaryKeyConstraint now
 function, tests added for before/after events on all constraint types. 
 [ticket:2105]

---
 CHANGES                   |  4 +++
 lib/sqlalchemy/schema.py  | 19 +++++++-------
 test/sql/test_metadata.py | 52 ++++++++++++++++++++++++++++++++++++---
 3 files changed, 61 insertions(+), 14 deletions(-)

diff --git a/CHANGES b/CHANGES
index 3aaf6f85f2..6db99a8ac2 100644
--- a/CHANGES
+++ b/CHANGES
@@ -100,6 +100,10 @@ CHANGES
     being placed between the "partition" and "order by"
     clauses.  [ticket:2134]
 
+  - Before/after attach events for PrimaryKeyConstraint
+    now function, tests added for before/after events
+    on all constraint types.  [ticket:2105]
+
 - engine
   - The C extension is now enabled by default on CPython
     2.x with a fallback to pure python if it fails to
diff --git a/lib/sqlalchemy/schema.py b/lib/sqlalchemy/schema.py
index af2c9f7721..45762490bc 100644
--- a/lib/sqlalchemy/schema.py
+++ b/lib/sqlalchemy/schema.py
@@ -285,7 +285,7 @@ class Table(SchemaItem, expression.TableClause):
         self.indexes = set()
         self.constraints = set()
         self._columns = expression.ColumnCollection()
-        self._set_primary_key(PrimaryKeyConstraint())
+        PrimaryKeyConstraint()._set_parent_with_dispatch(self) 
         self.foreign_keys = set()
         self._extra_dependencies = set()
         self.kwargs = {}
@@ -378,14 +378,6 @@ class Table(SchemaItem, expression.TableClause):
     def _init_collections(self):
         pass
 
-    def _set_primary_key(self, pk):
-        if self.primary_key in self.constraints:
-            self.constraints.remove(self.primary_key)
-        self.primary_key = pk
-        self.constraints.add(pk)
-
-        for c in pk.columns:
-            c.primary_key = True
 
     @util.memoized_property
     def _autoincrement_column(self):
@@ -2034,7 +2026,14 @@ class PrimaryKeyConstraint(ColumnCollectionConstraint):
 
     def _set_parent(self, table):
         super(PrimaryKeyConstraint, self)._set_parent(table)
-        table._set_primary_key(self)
+
+        if table.primary_key in table.constraints:
+            table.constraints.remove(table.primary_key)
+        table.primary_key = self
+        table.constraints.add(self)
+
+        for c in self.columns:
+            c.primary_key = True
 
     def _replace(self, col):
         self.columns.replace(col)
diff --git a/test/sql/test_metadata.py b/test/sql/test_metadata.py
index de86580fd5..7ca31c8544 100644
--- a/test/sql/test_metadata.py
+++ b/test/sql/test_metadata.py
@@ -1007,8 +1007,8 @@ class CatchAllEventsTest(fixtures.TestBase):
         def after_attach(obj, parent):
             canary.append("%s->%s" % (obj, parent))
 
-        event.listen(events.SchemaEventTarget, "before_parent_attach", before_attach)
-        event.listen(events.SchemaEventTarget, "after_parent_attach", after_attach)
+        event.listen(schema.SchemaItem, "before_parent_attach", before_attach)
+        event.listen(schema.SchemaItem, "after_parent_attach", after_attach)
 
         m = MetaData()
         t1 = Table('t1', m, 
@@ -1019,7 +1019,6 @@ class CatchAllEventsTest(fixtures.TestBase):
             Column('id', Integer, primary_key=True),
         )
 
-        # TODO: test more conditions here, constraints, defaults, etc.
         eq_(
             canary,
             [
@@ -1028,12 +1027,57 @@ class CatchAllEventsTest(fixtures.TestBase):
                 'ForeignKey->Column', 
                 "ForeignKey('t2.id')->bar", 
                 'Table->MetaData', 
+                'PrimaryKeyConstraint->Table', 'PrimaryKeyConstraint()->t1',
                 'Column->Table', 't1.id->t1', 
                 'Column->Table', 't1.bar->t1', 
                 'ForeignKeyConstraint->Table', 
                 'ForeignKeyConstraint()->t1', 
                 't1->MetaData(None)', 
-                'Table->MetaData', 'Column->Table', 
+                'Table->MetaData', 
+                'PrimaryKeyConstraint->Table', 'PrimaryKeyConstraint()->t2', 
+                'Column->Table', 
                 't2.id->t2', 't2->MetaData(None)']
         )
 
+    def test_events_per_constraint(self):
+        canary = []
+
+        def evt(target):
+            def before_attach(obj, parent):
+                canary.append("%s->%s" % (target.__name__, parent.__class__.__name__))
+
+            def after_attach(obj, parent):
+                canary.append("%s->%s" % (target.__name__, parent))
+            event.listen(target, "before_parent_attach", before_attach)
+            event.listen(target, "after_parent_attach", after_attach)
+
+        for target in [
+            schema.ForeignKeyConstraint, schema.PrimaryKeyConstraint, schema.UniqueConstraint,
+            schema.CheckConstraint
+        ]:
+            evt(target)
+
+        m = MetaData()
+        t1 = Table('t1', m, 
+            Column('id', Integer, Sequence('foo_id'), primary_key=True),
+            Column('bar', String, ForeignKey('t2.id')),
+            Column('bat', Integer, unique=True),
+        )
+        t2 = Table('t2', m,
+            Column('id', Integer, primary_key=True),
+            Column('bar', Integer),
+            Column('bat', Integer),
+            CheckConstraint("bar>5"),
+            UniqueConstraint('bar', 'bat')
+        )
+        eq_(
+            canary,
+            [
+            'PrimaryKeyConstraint->Table', 'PrimaryKeyConstraint->t1', 
+            'ForeignKeyConstraint->Table', 'ForeignKeyConstraint->t1',
+            'UniqueConstraint->Table', 'UniqueConstraint->t1',
+            'PrimaryKeyConstraint->Table', 'PrimaryKeyConstraint->t2', 
+            'CheckConstraint->Table', 'CheckConstraint->t2',
+            'UniqueConstraint->Table', 'UniqueConstraint->t2'
+            ]
+        )
-- 
2.39.5