from sqlalchemy import schema as sa_schema, types as sqltypes, sql import logging from .. import compat import re from ..compat import string_types log = logging.getLogger(__name__) try: from sqlalchemy.sql.naming import conv def _render_gen_name(autogen_context, name): if isinstance(name, conv): return _f_name(_alembic_autogenerate_prefix(autogen_context), name) else: return name except ImportError: def _render_gen_name(autogen_context, name): return name class _f_name(object): def __init__(self, prefix, name): self.prefix = prefix self.name = name def __repr__(self): return "%sf(%r)" % (self.prefix, self.name) def _render_potential_expr(value, autogen_context): if isinstance(value, sql.ClauseElement): if compat.sqla_08: compile_kw = dict(compile_kwargs={'literal_binds': True}) else: compile_kw = {} return "%(prefix)stext(%(sql)r)" % { "prefix": _sqlalchemy_autogenerate_prefix(autogen_context), "sql": str( value.compile(dialect=autogen_context['dialect'], **compile_kw) ) } else: return repr(value) def _add_table(table, autogen_context): text = "%(prefix)screate_table(%(tablename)r,\n%(args)s" % { 'tablename': table.name, 'prefix': _alembic_autogenerate_prefix(autogen_context), 'args': ',\n'.join( [col for col in [_render_column(col, autogen_context) for col in table.c] if col] + sorted([rcons for rcons in [_render_constraint(cons, autogen_context) for cons in table.constraints] if rcons is not None ]) ) } if table.schema: text += ",\nschema=%r" % table.schema for k in sorted(table.kwargs): text += ",\n%s=%r" % (k.replace(" ", "_"), table.kwargs[k]) text += "\n)" return text def _drop_table(table, autogen_context): text = "%(prefix)sdrop_table(%(tname)r" % { "prefix": _alembic_autogenerate_prefix(autogen_context), "tname": table.name } if table.schema: text += ", schema=%r" % table.schema text += ")" return text def _add_index(index, autogen_context): """ Generate Alembic operations for the CREATE INDEX of an :class:`~sqlalchemy.schema.Index` instance. """ from .compare import _get_index_column_names text = "%(prefix)screate_index('%(name)s', '%(table)s', %(columns)s, "\ "unique=%(unique)r%(schema)s%(kwargs)s)" % { 'prefix': _alembic_autogenerate_prefix(autogen_context), 'name': _render_gen_name(autogen_context, index.name), 'table': index.table.name, 'columns': _get_index_column_names(index), 'unique': index.unique or False, 'schema': (", schema='%s'" % index.table.schema) if index.table.schema else '', 'kwargs': (', '+', '.join( ["%s=%s" % (key, _render_potential_expr(val, autogen_context)) for key, val in index.kwargs.items()]))\ if len(index.kwargs) else '' } return text def _drop_index(index, autogen_context): """ Generate Alembic operations for the DROP INDEX of an :class:`~sqlalchemy.schema.Index` instance. """ text = "%(prefix)sdrop_index('%(name)s', "\ "table_name='%(table_name)s'%(schema)s)" % { 'prefix': _alembic_autogenerate_prefix(autogen_context), 'name': _render_gen_name(autogen_context, index.name), 'table_name': index.table.name, 'schema': ((", schema='%s'" % index.table.schema) if index.table.schema else '') } return text def _render_unique_constraint(constraint, autogen_context): rendered = _user_defined_render("unique", constraint, autogen_context) if rendered is not False: return rendered return _uq_constraint(constraint, autogen_context, False) def _add_unique_constraint(constraint, autogen_context): """ Generate Alembic operations for the ALTER TABLE .. ADD CONSTRAINT ... UNIQUE of a :class:`~sqlalchemy.schema.UniqueConstraint` instance. """ return _uq_constraint(constraint, autogen_context, True) def _uq_constraint(constraint, autogen_context, alter): opts = [] if constraint.deferrable: opts.append(("deferrable", str(constraint.deferrable))) if constraint.initially: opts.append(("initially", str(constraint.initially))) if alter and constraint.table.schema: opts.append(("schema", str(constraint.table.schema))) if not alter and constraint.name: opts.append(("name", _render_gen_name(autogen_context, constraint.name))) if alter: args = [repr(_render_gen_name(autogen_context, constraint.name)), repr(constraint.table.name)] args.append(repr([col.name for col in constraint.columns])) args.extend(["%s=%r" % (k, v) for k, v in opts]) return "%(prefix)screate_unique_constraint(%(args)s)" % { 'prefix': _alembic_autogenerate_prefix(autogen_context), 'args': ", ".join(args) } else: args = [repr(col.name) for col in constraint.columns] args.extend(["%s=%r" % (k, v) for k, v in opts]) return "%(prefix)sUniqueConstraint(%(args)s)" % { "prefix": _sqlalchemy_autogenerate_prefix(autogen_context), "args": ", ".join(args) } def _add_fk_constraint(constraint, autogen_context): raise NotImplementedError() def _add_pk_constraint(constraint, autogen_context): raise NotImplementedError() def _add_check_constraint(constraint, autogen_context): raise NotImplementedError() def _add_constraint(constraint, autogen_context): """ Dispatcher for the different types of constraints. """ funcs = { "unique_constraint": _add_unique_constraint, "foreign_key_constraint": _add_fk_constraint, "primary_key_constraint": _add_pk_constraint, "check_constraint": _add_check_constraint, "column_check_constraint": _add_check_constraint, } return funcs[constraint.__visit_name__](constraint, autogen_context) def _drop_constraint(constraint, autogen_context): """ Generate Alembic operations for the ALTER TABLE ... DROP CONSTRAINT of a :class:`~sqlalchemy.schema.UniqueConstraint` instance. """ text = "%(prefix)sdrop_constraint(%(name)r, '%(table_name)s'%(schema)s)" % { 'prefix': _alembic_autogenerate_prefix(autogen_context), 'name': _render_gen_name(autogen_context, constraint.name), 'table_name': constraint.table.name, 'schema': (", schema='%s'" % constraint.table.schema) if constraint.table.schema else '', } return text def _add_column(schema, tname, column, autogen_context): text = "%(prefix)sadd_column(%(tname)r, %(column)s" % { "prefix": _alembic_autogenerate_prefix(autogen_context), "tname": tname, "column": _render_column(column, autogen_context) } if schema: text += ", schema=%r" % schema text += ")" return text def _drop_column(schema, tname, column, autogen_context): text = "%(prefix)sdrop_column(%(tname)r, %(cname)r" % { "prefix": _alembic_autogenerate_prefix(autogen_context), "tname": tname, "cname": column.name } if schema: text += ", schema=%r" % schema text += ")" return text def _modify_col(tname, cname, autogen_context, server_default=False, type_=None, nullable=None, existing_type=None, existing_nullable=None, existing_server_default=False, schema=None): indent = " " * 11 text = "%(prefix)salter_column(%(tname)r, %(cname)r" % { 'prefix': _alembic_autogenerate_prefix( autogen_context), 'tname': tname, 'cname': cname} text += ",\n%sexisting_type=%s" % (indent, _repr_type(existing_type, autogen_context)) if server_default is not False: rendered = _render_server_default( server_default, autogen_context) text += ",\n%sserver_default=%s" % (indent, rendered) if type_ is not None: text += ",\n%stype_=%s" % (indent, _repr_type(type_, autogen_context)) if nullable is not None: text += ",\n%snullable=%r" % ( indent, nullable,) if existing_nullable is not None: text += ",\n%sexisting_nullable=%r" % ( indent, existing_nullable) if existing_server_default: rendered = _render_server_default( existing_server_default, autogen_context) text += ",\n%sexisting_server_default=%s" % ( indent, rendered) if schema: text += ",\n%sschema=%r" % (indent, schema) text += ")" return text def _user_autogenerate_prefix(autogen_context): prefix = autogen_context['opts']['user_module_prefix'] if prefix is None: return _sqlalchemy_autogenerate_prefix(autogen_context) else: return prefix def _sqlalchemy_autogenerate_prefix(autogen_context): return autogen_context['opts']['sqlalchemy_module_prefix'] or '' def _alembic_autogenerate_prefix(autogen_context): return autogen_context['opts']['alembic_module_prefix'] or '' def _user_defined_render(type_, object_, autogen_context): if 'opts' in autogen_context and \ 'render_item' in autogen_context['opts']: render = autogen_context['opts']['render_item'] if render: rendered = render(type_, object_, autogen_context) if rendered is not False: return rendered return False def _render_column(column, autogen_context): rendered = _user_defined_render("column", column, autogen_context) if rendered is not False: return rendered opts = [] if column.server_default: rendered = _render_server_default( column.server_default, autogen_context ) if rendered: opts.append(("server_default", rendered)) if not column.autoincrement: opts.append(("autoincrement", column.autoincrement)) if column.nullable is not None: opts.append(("nullable", column.nullable)) # TODO: for non-ascii colname, assign a "key" return "%(prefix)sColumn(%(name)r, %(type)s, %(kw)s)" % { 'prefix': _sqlalchemy_autogenerate_prefix(autogen_context), 'name': column.name, 'type': _repr_type(column.type, autogen_context), 'kw': ", ".join(["%s=%s" % (kwname, val) for kwname, val in opts]) } def _render_server_default(default, autogen_context): rendered = _user_defined_render("server_default", default, autogen_context) if rendered is not False: return rendered if isinstance(default, sa_schema.DefaultClause): if isinstance(default.arg, string_types): default = default.arg else: default = str(default.arg.compile( dialect=autogen_context['dialect'])) if isinstance(default, string_types): # TODO: this is just a hack to get # tests to pass until we figure out # WTF sqlite is doing default = re.sub(r"^'|'$", "", default) return repr(default) else: return None def _repr_type(type_, autogen_context): rendered = _user_defined_render("type", type_, autogen_context) if rendered is not False: return rendered mod = type(type_).__module__ imports = autogen_context.get('imports', None) if mod.startswith("sqlalchemy.dialects"): dname = re.match(r"sqlalchemy\.dialects\.(\w+)", mod).group(1) if imports is not None: imports.add("from sqlalchemy.dialects import %s" % dname) return "%s.%r" % (dname, type_) elif mod.startswith("sqlalchemy"): prefix = _sqlalchemy_autogenerate_prefix(autogen_context) return "%s%r" % (prefix, type_) else: prefix = _user_autogenerate_prefix(autogen_context) return "%s%r" % (prefix, type_) def _render_constraint(constraint, autogen_context): renderer = _constraint_renderers.get(type(constraint), None) if renderer: return renderer(constraint, autogen_context) else: return None def _render_primary_key(constraint, autogen_context): rendered = _user_defined_render("primary_key", constraint, autogen_context) if rendered is not False: return rendered if not constraint.columns: return None opts = [] if constraint.name: opts.append(("name", repr(_render_gen_name(autogen_context, constraint.name)))) return "%(prefix)sPrimaryKeyConstraint(%(args)s)" % { "prefix": _sqlalchemy_autogenerate_prefix(autogen_context), "args": ", ".join( [repr(c.key) for c in constraint.columns] + ["%s=%s" % (kwname, val) for kwname, val in opts] ), } def _fk_colspec(fk, metadata_schema): """Implement a 'safe' version of ForeignKey._get_colspec() that never tries to resolve the remote table. """ if metadata_schema is None: return fk._get_colspec() else: # need to render schema breaking up tokens by hand, since the # ForeignKeyConstraint here may not actually have a remote # Table present tokens = fk._colspec.split(".") # no schema in the colspec, render it if len(tokens) == 2: return "%s.%s" % (metadata_schema, fk._colspec) else: return fk._colspec def _render_foreign_key(constraint, autogen_context): rendered = _user_defined_render("foreign_key", constraint, autogen_context) if rendered is not False: return rendered opts = [] if constraint.name: opts.append(("name", repr(_render_gen_name(autogen_context, constraint.name)))) if constraint.onupdate: opts.append(("onupdate", repr(constraint.onupdate))) if constraint.ondelete: opts.append(("ondelete", repr(constraint.ondelete))) if constraint.initially: opts.append(("initially", repr(constraint.initially))) if constraint.deferrable: opts.append(("deferrable", repr(constraint.deferrable))) if constraint.use_alter: opts.append(("use_alter", repr(constraint.use_alter))) apply_metadata_schema = constraint.parent.metadata.schema return "%(prefix)sForeignKeyConstraint([%(cols)s], "\ "[%(refcols)s], %(args)s)" % { "prefix": _sqlalchemy_autogenerate_prefix(autogen_context), "cols": ", ".join("'%s'" % f.parent.key for f in constraint.elements), "refcols": ", ".join(repr(_fk_colspec(f, apply_metadata_schema)) for f in constraint.elements), "args": ", ".join( ["%s=%s" % (kwname, val) for kwname, val in opts] ), } def _render_check_constraint(constraint, autogen_context): rendered = _user_defined_render("check", constraint, autogen_context) if rendered is not False: return rendered # detect the constraint being part of # a parent type which is probably in the Table already. # ideally SQLAlchemy would give us more of a first class # way to detect this. if constraint._create_rule and \ hasattr(constraint._create_rule, 'target') and \ isinstance(constraint._create_rule.target, sqltypes.TypeEngine): return None opts = [] if constraint.name: opts.append(("name", repr(_render_gen_name(autogen_context, constraint.name)))) return "%(prefix)sCheckConstraint(%(sqltext)r%(opts)s)" % { "prefix": _sqlalchemy_autogenerate_prefix(autogen_context), "opts": ", " + (", ".join("%s=%s" % (k, v) for k, v in opts)) if opts else "", "sqltext": str( constraint.sqltext.compile( dialect=autogen_context['dialect'] ) ) } _constraint_renderers = { sa_schema.PrimaryKeyConstraint: _render_primary_key, sa_schema.ForeignKeyConstraint: _render_foreign_key, sa_schema.UniqueConstraint: _render_unique_constraint, sa_schema.CheckConstraint: _render_check_constraint }