From: Mike Bayer <mike_mp@zzzcomputing.com>
Date: Mon, 26 Apr 2010 04:20:57 +0000 (-0400)
Subject: illustrate the basic idea of add_xxx(strings), whereby Table/Column/Constraint objects
X-Git-Tag: rel_0_1_0~109
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=de3643bc4fb5c6af94fb2139395950effb09d281;p=thirdparty%2Fsqlalchemy%2Falembic.git

illustrate the basic idea of add_xxx(strings), whereby Table/Column/Constraint objects
are just generated in order to produce the DDL.   I'm not entirely
thrilled with this - would be nicer if we could get sqlalchemy.schema constructs
to generate with strings alone - but this would change their signatures.  so we're sort
of doing that here anyway.
---

diff --git a/alembic/command.py b/alembic/command.py
index 5110db26..9e79c070 100644
--- a/alembic/command.py
+++ b/alembic/command.py
@@ -14,7 +14,7 @@ def list_templates(opts):
                         tempname, 
                         'README')
         synopsis = open(readme).next()
-        print options.format_opt(tempname, synopsis)
+        print util.format_opt(tempname, synopsis)
     
     print "\nTemplates are used via the 'init' command, e.g.:"
     print "\n  alembic init --template pylons ./scripts"
diff --git a/alembic/context.py b/alembic/context.py
index 27771daa..a919c07b 100644
--- a/alembic/context.py
+++ b/alembic/context.py
@@ -1,4 +1,5 @@
 from alembic.ddl import base
+from alembic import util
 
 class ContextMeta(type):
     def __init__(cls, classname, bases, dict_):
@@ -15,17 +16,25 @@ class DefaultContext(object):
     def __init__(self, options, connection):
         self.options = options
         self.connection = connection
+    
+    def _exec(self, construct):
+        pass
         
     def alter_column(self, table_name, column_name, 
-                        nullable=NO_VALUE,
-                        server_default=NO_VALUE,
-                        name=NO_VALUE,
-                        type=NO_VALUE
+                        nullable=util.NO_VALUE,
+                        server_default=util.NO_VALUE,
+                        name=util.NO_VALUE,
+                        type=util.NO_VALUE
     ):
     
-        if nullable is not NO_VALUE:
-            base.ColumnNullable(table_name, column_name, nullable)
-        if server_default is not NO_VALUE:
-            base.ColumnDefault(table_name, column_name, server_default)
+        if nullable is not util.NO_VALUE:
+            self._exec(base.ColumnNullable(table_name, column_name, nullable))
+        if server_default is not util.NO_VALUE:
+            self._exec(base.ColumnDefault(table_name, column_name, server_default))
     
-        # ... etc
\ No newline at end of file
+        # ... etc
+        
+    def add_constraint(self, const):
+        self._exec(schema.AddConstraint(const))
+
+        
\ No newline at end of file
diff --git a/alembic/op.py b/alembic/op.py
index 6a31c4f0..16fc1ab7 100644
--- a/alembic/op.py
+++ b/alembic/op.py
@@ -1,12 +1,14 @@
-from sqlalchemy import util
+from alembic import util
+from sqlalchemy.types import NULLTYPE
+from sqlalchemy import schema
 
-NO_VALUE = util.symbol("NO_VALUE")
+__all__ = ['alter_column', 'add_foreign_key']
 
 def alter_column(table_name, column_name, 
-                    nullable=NO_VALUE,
-                    server_default=NO_VALUE,
-                    name=NO_VALUE,
-                    type_=NO_VALUE
+                    nullable=util.NO_VALUE,
+                    server_default=util.NO_VALUE,
+                    name=util.NO_VALUE,
+                    type_=util.NO_VALUE
 ):
     """Issue ALTER COLUMN using the current change context."""
     
@@ -16,3 +18,34 @@ def alter_column(table_name, column_name,
         name=name,
         type_=type_
     )
+
+
+def _foreign_key_constraint(name, source, referent, local_cols, remote_cols):
+    m = schema.MetaData()
+    t1 = schema.Table(source, m, 
+            *[schema.Column(n, NULLTYPE) for n in local_cols])
+    t2 = schema.Table(referent, m, 
+            *[schema.Column(n, NULLTYPE) for n in remote_cols])
+
+    f = schema.ForeignKeyConstraint(local_cols, 
+                                        ["%s.%s" % (referent, name) 
+                                        for name in remote_cols],
+                                        name=name
+                                        )
+    t1.append_constraint(f)
+    return f
+
+def _unique_constraint(name, source, local_cols):
+    t = schema.Table(source, schema.MetaData(), 
+                *[schema.Column(n, NULLTYPE) for n in local_cols])
+    return schema.UniqueConstraint(*t.c, name=name)
+    
+def create_foreign_key(name, source, referent, local_cols, remote_cols):
+    context.add_constraint(
+                _foreign_key_constraint(source, referent, local_cols, remote_cols)
+            )
+
+def create_unique_constraint(name, source, local_cols):
+    context.add_constraint(
+                _unique_constraint(name, source, local_cols)
+            )
diff --git a/alembic/util.py b/alembic/util.py
index a7ae8852..54004bdb 100644
--- a/alembic/util.py
+++ b/alembic/util.py
@@ -2,6 +2,10 @@ from mako.template import Template
 import sys
 import os
 import textwrap
+from sqlalchemy import util
+
+NO_VALUE = util.symbol("NO_VALUE")
+
 
 def template_to_file(template_file, dest, **kw):
     f = open(dest, 'w')
diff --git a/tests/test_schema.py b/tests/test_schema.py
new file mode 100644
index 00000000..5ba4fc9e
--- /dev/null
+++ b/tests/test_schema.py
@@ -0,0 +1,18 @@
+from tests import assert_compiled
+from alembic import op
+from sqlalchemy.schema import AddConstraint, ForeignKeyConstraint
+
+def test_foreign_key():
+    fk = op._foreign_key_constraint('fk_test', 't1', 't2', ['foo', 'bar'], ['bat', 'hoho'])
+    assert_compiled(
+        AddConstraint(fk),
+        "ALTER TABLE t1 ADD CONSTRAINT hoho FOREIGN KEY(foo, bar) REFERENCES t2 (bat, hoho)"
+    )
+    
+def test_unique_constraint():
+    uc = op._unique_constraint('uk_test', 't1', ['foo', 'bar'])
+    assert_compiled(
+        AddConstraint(uc),
+        "ALTER TABLE t1 ADD CONSTRAINT uk_test UNIQUE (foo, bar)"
+    )
+    
\ No newline at end of file