IDENTITY column attribute

Author: oyiptong

redshift_sqlalchemy/dialect.py

@@ -4,7 +4,7 @@
 from sqlalchemy.engine import reflection
 from sqlalchemy.ext.compiler import compiles
 from sqlalchemy.sql.expression import BindParameter, Executable, ClauseElement
-from sqlalchemy.types import VARCHAR, NullType
+from sqlalchemy.types import VARCHAR, NullType, BigInteger, Integer
 
 
 class RedShiftDDLCompiler(PGDDLCompiler):
@@ -75,17 +75,16 @@ def post_create_table(self, table):
         return text
 
     def get_column_specification(self, column, **kwargs):
-        # aron - Apr 21, 2014: Redshift doesn't support serial types. So I
-        # removed support for them here.
         colspec = self.preparer.format_column(column)
-        colspec += " " + self.dialect.type_compiler.process(column.type)
 
-        colspec += self._fetch_redshift_column_attributes(column)
+        colspec += " " + self.dialect.type_compiler.process(column.type)
 
         default = self.get_column_default_string(column)
         if default is not None:
             colspec += " DEFAULT " + default
 
+        colspec += self._fetch_redshift_column_attributes(column)
+
         if not column.nullable:
             colspec += " NOT NULL"
         return colspec
@@ -95,6 +94,10 @@ def _fetch_redshift_column_attributes(self, column):
         if not hasattr(column, 'info'):
             return text
         info = column.info
+        identity = info.get('identity', None)
+        if identity:
+            text += " IDENTITY({0},{1})".format(identity[0], identity[1])
+
         encode = info.get('encode', None)
         if encode:
             text += " ENCODE " + encode

tests/test_ddl_compiler.py

@@ -1,7 +1,7 @@
 import difflib
 
 from pytest import fixture
-from sqlalchemy import Table, Column, Integer, String, MetaData
+from sqlalchemy import Table, Column, Integer, String, Sequence, MetaData
 from sqlalchemy.exc import CompileError
 from sqlalchemy.schema import CreateTable
 
@@ -40,6 +40,22 @@ def test_create_table_simple(self, compiler):
                     u"\n\tPRIMARY KEY (id)\n)\n\n"
         assert expected == actual, self._compare_strings(expected, actual)
 
+    def test_create_table_with_identity(self, compiler):
+
+        table = Table('t1',
+                      MetaData(),
+                      Column('id', Integer, primary_key=True, info={'identity': [1, 2]}),
+                      Column('name', String))
+
+
+        create_table = CreateTable(table)
+        actual = compiler.process(create_table)
+        expected =  u"\nCREATE TABLE t1 ("\
+                    u"\n\tid INTEGER IDENTITY(1,2) NOT NULL, "\
+                    u"\n\tname VARCHAR, "\
+                    u"\n\tPRIMARY KEY (id)\n)\n\n"
+        assert expected == actual, self._compare_strings(expected, actual)
+
     def test_create_table_with_diststyle(self, compiler):
 
         table = Table('t1',