Merge "Add support for identity columns"

2 files changed, 198 insertions, 23 deletions

diff --git a/lib/sqlalchemy/sql/compiler.py b/lib/sqlalchemy/sql/compiler.py
index 4f4cf7f8b..17cacc981 100644
--- a/lib/sqlalchemy/sql/compiler.py
+++ b/lib/sqlalchemy/sql/compiler.py
@@ -3918,31 +3918,39 @@ class DDLCompiler(Compiled):
             drop.element, use_table=True
         )
 
+    def get_identity_options(self, identity_options):
+        text = []
+        if identity_options.increment is not None:
+            text.append("INCREMENT BY %d" % identity_options.increment)
+        if identity_options.start is not None:
+            text.append("START WITH %d" % identity_options.start)
+        if identity_options.minvalue is not None:
+            text.append("MINVALUE %d" % identity_options.minvalue)
+        if identity_options.maxvalue is not None:
+            text.append("MAXVALUE %d" % identity_options.maxvalue)
+        if identity_options.nominvalue is not None:
+            text.append("NO MINVALUE")
+        if identity_options.nomaxvalue is not None:
+            text.append("NO MAXVALUE")
+        if identity_options.cache is not None:
+            text.append("CACHE %d" % identity_options.cache)
+        if identity_options.order is True:
+            text.append("ORDER")
+        if identity_options.cycle is not None:
+            text.append("CYCLE")
+        return " ".join(text)
+
     def visit_create_sequence(self, create, prefix=None, **kw):
         text = "CREATE SEQUENCE %s" % self.preparer.format_sequence(
             create.element
         )
         if prefix:
             text += prefix
-        if create.element.increment is not None:
-            text += " INCREMENT BY %d" % create.element.increment
         if create.element.start is None:
             create.element.start = self.dialect.default_sequence_base
-        text += " START WITH %d" % create.element.start
-        if create.element.minvalue is not None:
-            text += " MINVALUE %d" % create.element.minvalue
-        if create.element.maxvalue is not None:
-            text += " MAXVALUE %d" % create.element.maxvalue
-        if create.element.nominvalue is not None:
-            text += " NO MINVALUE"
-        if create.element.nomaxvalue is not None:
-            text += " NO MAXVALUE"
-        if create.element.cache is not None:
-            text += " CACHE %d" % create.element.cache
-        if create.element.order is True:
-            text += " ORDER"
-        if create.element.cycle is not None:
-            text += " CYCLE"
+        options = self.get_identity_options(create.element)
+        if options:
+            text += " " + options
         return text
 
     def visit_drop_sequence(self, drop, **kw):
@@ -3981,6 +3989,9 @@ class DDLCompiler(Compiled):
         if column.computed is not None:
             colspec += " " + self.process(column.computed)
 
+        if column.identity is not None:
+            colspec += " " + self.process(column.identity)
+
         if not column.nullable:
             colspec += " NOT NULL"
         return colspec
@@ -4138,6 +4149,15 @@ class DDLCompiler(Compiled):
             text += " VIRTUAL"
         return text
 
+    def visit_identity_column(self, identity, **kw):
+        text = "GENERATED %s AS IDENTITY" % (
+            "ALWAYS" if identity.always else "BY DEFAULT",
+        )
+        options = self.get_identity_options(identity)
+        if options:
+            text += " (%s)" % options
+        return text
+
 
 class GenericTypeCompiler(TypeCompiler):
     def visit_FLOAT(self, type_, **kw):
diff --git a/lib/sqlalchemy/sql/schema.py b/lib/sqlalchemy/sql/schema.py
index 0b04ff0da..c2a41205f 100644
--- a/lib/sqlalchemy/sql/schema.py
+++ b/lib/sqlalchemy/sql/schema.py
@@ -1098,8 +1098,8 @@ class Column(DialectKWArgs, SchemaItem, ColumnClause):
           :class:`.SchemaItem` derived constructs which will be applied
           as options to the column.  These include instances of
           :class:`.Constraint`, :class:`_schema.ForeignKey`,
-          :class:`.ColumnDefault`,
-          :class:`.Sequence`, :class:`.Computed`.  In some cases an
+          :class:`.ColumnDefault`, :class:`.Sequence`, :class:`.Computed`
+          :class:`.Identity`.  In some cases an
           equivalent keyword argument is available such as ``server_default``,
           ``default`` and ``unique``.
 
@@ -1113,7 +1113,9 @@ class Column(DialectKWArgs, SchemaItem, ColumnClause):
           AUTO_INCREMENT will be emitted for this column during a table
           create, as well as that the column is assumed to generate new
           integer primary key values when an INSERT statement invokes which
-          will be retrieved by the dialect.
+          will be retrieved by the dialect.  When used in conjunction with
+          :class:`.Identity` on a dialect that supports it, this parameter
+          has no effect.
 
           The flag may be set to ``True`` to indicate that a column which
           is part of a composite (e.g. multi-column) primary key should
@@ -1381,6 +1383,7 @@ class Column(DialectKWArgs, SchemaItem, ColumnClause):
         self.foreign_keys = set()
         self.comment = kwargs.pop("comment", None)
         self.computed = None
+        self.identity = None
 
         # check if this Column is proxying another column
         if "_proxies" in kwargs:
@@ -1563,6 +1566,14 @@ class Column(DialectKWArgs, SchemaItem, ColumnClause):
 
         self._setup_on_memoized_fks(lambda fk: fk._set_remote_table(table))
 
+        if self.identity and (
+            isinstance(self.default, Sequence)
+            or isinstance(self.onupdate, Sequence)
+        ):
+            raise exc.ArgumentError(
+                "An column cannot specify both Identity and Sequence."
+            )
+
     def _setup_on_memoized_fks(self, fn):
         fk_keys = [
             ((self.table.key, self.key), False),
@@ -1606,7 +1617,7 @@ class Column(DialectKWArgs, SchemaItem, ColumnClause):
 
         server_default = self.server_default
         server_onupdate = self.server_onupdate
-        if isinstance(server_default, Computed):
+        if isinstance(server_default, (Computed, Identity)):
             server_default = server_onupdate = None
             args.append(self.server_default.copy(**kw))
 
@@ -2369,7 +2380,7 @@ class IdentityOptions(object):
         """Construct a :class:`.IdentityOptions` object.
 
         See the :class:`.Sequence` documentation for a complete description
-        of the parameters
+        of the parameters.
 
         :param start: the starting index of the sequence.
         :param increment: the increment value of the sequence.
@@ -3602,7 +3613,11 @@ class PrimaryKeyConstraint(ColumnCollectionConstraint):
                 and not autoinc_true
             ):
                 return False
-            elif col.server_default is not None and not autoinc_true:
+            elif (
+                col.server_default is not None
+                and not isinstance(col.server_default, Identity)
+                and not autoinc_true
+            ):
                 return False
             elif col.foreign_keys and col.autoincrement not in (
                 True,
@@ -4612,3 +4627,143 @@ class Computed(FetchedValue, SchemaItem):
         g = Computed(sqltext, persisted=self.persisted)
 
         return self._schema_item_copy(g)
+
+
+class Identity(IdentityOptions, FetchedValue, SchemaItem):
+    """Defines an identity column, i.e. "GENERATED { ALWAYS | BY DEFAULT }
+    AS IDENTITY" syntax.
+
+    The :class:`.Identity` construct is an inline construct added to the
+    argument list of a :class:`_schema.Column` object::
+
+        from sqlalchemy import Identity
+
+        Table('foo', meta,
+            Column('id', Integer, Identity())
+            Column('description', Text),
+        )
+
+    See the linked documentation below for complete details.
+
+    .. versionadded:: 1.4
+
+    .. seealso::
+
+        :ref:`identity_ddl`
+
+    """
+
+    __visit_name__ = "identity_column"
+
+    def __init__(
+        self,
+        always=False,
+        on_null=None,
+        start=None,
+        increment=None,
+        minvalue=None,
+        maxvalue=None,
+        nominvalue=None,
+        nomaxvalue=None,
+        cycle=None,
+        cache=None,
+        order=None,
+    ):
+        """Construct a GENERATED { ALWAYS | BY DEFAULT } AS IDENTITY DDL
+        construct to accompany a :class:`_schema.Column`.
+
+        See the :class:`.Sequence` documentation for a complete description
+        of most parameters.
+
+        .. note::
+            MSSQL supports this construct as the preferred alternative to
+            generate an IDENTITY on a column, but it uses non standard
+            syntax that only support :paramref:`_schema.Identity.start`
+            and :paramref:`_schema.Identity.increment`.
+            All other parameters are ignored.
+
+        :param always:
+          A boolean, that indicates the type of identity column.
+          If ``False`` is specified, the default, then the user-specified
+          value takes precedence.
+          If ``True`` is specified, a user-specified value is not accepted (
+          on some backends, like PostgreSQL, OVERRIDING SYSTEM VALUE, or
+          similar, may be specified in an INSERT to override the sequence
+          value).
+          Some backends also have a default value for this parameter,
+          ``None`` can be used to omit rendering this part in the DDL. It
+          will be treated as ``False`` if a backend does not have a default
+          value.
+
+        :param on_null:
+          Set to ``True`` to specify ON NULL in conjunction with a
+          'by default' identity column. This option is only supported on
+          some backends, like Oracle.
+
+        :param start: the starting index of the sequence.
+        :param increment: the increment value of the sequence.
+        :param minvalue: the minimum value of the sequence.
+        :param maxvalue: the maximum value of the sequence.
+        :param nominvalue: no minimum value of the sequence.
+        :param nomaxvalue: no maximum value of the sequence.
+        :param cycle: allows the sequence to wrap around when the maxvalue
+         or minvalue has been reached.
+        :param cache: optional integer value; number of future values in the
+         sequence which are calculated in advance.
+        :param order: optional boolean value; if true, renders the
+         ORDER keyword.
+
+        """
+        IdentityOptions.__init__(
+            self,
+            start=start,
+            increment=increment,
+            minvalue=minvalue,
+            maxvalue=maxvalue,
+            nominvalue=nominvalue,
+            nomaxvalue=nomaxvalue,
+            cycle=cycle,
+            cache=cache,
+            order=order,
+        )
+        self.always = always
+        self.on_null = on_null
+        self.column = None
+
+    def _set_parent(self, parent):
+        if not isinstance(
+            parent.server_default, (type(None), Identity)
+        ) or not isinstance(parent.server_onupdate, type(None)):
+            raise exc.ArgumentError(
+                "A column with an Identity object cannot specify a "
+                "server_default or a server_onupdate argument"
+            )
+        if parent.autoincrement is False:
+            raise exc.ArgumentError(
+                "A column with an Identity object cannot specify "
+                "autoincrement=False"
+            )
+        self.column = parent
+        parent.identity = self
+        # self.column.server_onupdate = self
+        self.column.server_default = self
+
+    def _as_for_update(self, for_update):
+        return self
+
+    def copy(self, target_table=None, **kw):
+        i = Identity(
+            always=self.always,
+            on_null=self.on_null,
+            start=self.start,
+            increment=self.increment,
+            minvalue=self.minvalue,
+            maxvalue=self.maxvalue,
+            nominvalue=self.nominvalue,
+            nomaxvalue=self.nomaxvalue,
+            cycle=self.cycle,
+            cache=self.cache,
+            order=self.order,
+        )
+
+        return self._schema_item_copy(i)