Add EnvironmentContext and Config to MigrationContext

[msabramo]

My use case is that in my migrations, I want to be able to execute SQL that uses a schema name and the schema name is something that will be configured differently depending on the user or CI or deployment system that is running.

I modified the code so that the MigrationContext has access to the EnvironmentContext and the config (which comes from `alembic.ini).

This lets migrations do stuff like:

# 13189a2dd36_populate_stuff.py

...
schema = op.get_context().config.get_main_option('schema')
...
def downgrade():
    print("downgrade (%r): schema = %r" % (revision, schema))
    op.execute("""
        TRUNCATE TABLE {schema}.FooTable;
        """.format(schema=schema))

where schema is a custom option that I added to alembic.ini

# alembic.ini

...
sqlalchemy.url = mssql+pymssql://...
schema = foobar_schema
...

[msabramo]

Oops. It seems calling op.get_context() in the global scope of a revision is not kosher.

❯ alembic revision -m "Populate sm_Bench_Source_D_Type"
Traceback (most recent call last):
  File "/Users/marca/dev/surveymonkey/profilesvc/bin/alembic", line 9, in <module>
    load_entry_point('alembic==0.6.5dev', 'console_scripts', 'alembic')()
  File "/Users/marca/dev/git-repos/alembic/alembic/config.py", line 298, in main
    CommandLine(prog=prog).main(argv=argv)
  File "/Users/marca/dev/git-repos/alembic/alembic/config.py", line 293, in main
    self.run_cmd(cfg, options)
  File "/Users/marca/dev/git-repos/alembic/alembic/config.py", line 279, in run_cmd
    **dict((k, getattr(options, k)) for k in kwarg)
  File "/Users/marca/dev/git-repos/alembic/alembic/command.py", line 99, in revision
    **template_args)
  File "/Users/marca/dev/git-repos/alembic/alembic/script.py", line 342, in generate_revision
    current_head = self.get_current_head()
  File "/Users/marca/dev/git-repos/alembic/alembic/script.py", line 259, in get_current_head
    current_heads = self.get_heads()
  File "/Users/marca/dev/git-repos/alembic/alembic/script.py", line 288, in get_heads
    for script in self._revision_map.values():
  File "/Users/marca/dev/git-repos/alembic/alembic/util.py", line 268, in __get__
    obj.__dict__[self.__name__] = result = self.fget(obj)
  File "/Users/marca/dev/git-repos/alembic/alembic/script.py", line 213, in _revision_map
    script = Script._from_filename(self, self.versions, file_)
  File "/Users/marca/dev/git-repos/alembic/alembic/script.py", line 496, in _from_filename
    module = util.load_python_file(dir_, filename)
  File "/Users/marca/dev/git-repos/alembic/alembic/util.py", line 212, in load_python_file
    module = load_module_py(module_id, path)
  File "/Users/marca/dev/git-repos/alembic/alembic/compat.py", line 58, in load_module_py
    mod = imp.load_source(module_id, path, fp)
  File "alembic/versions/13189a2dd36_populate_sm_bench_entity_d_type.py", line 19, in <module>
    schema = op.get_context().config.get_main_option('schema')
  File "<string>", line 6, in get_context
  File "/Users/marca/dev/git-repos/alembic/alembic/util.py", line 108, in _name_error
    name, cls.__name__
NameError: Can't invoke function 'get_context', as the proxy object has not yet been established for the Alembic 'Operations' class.  Try placing this code inside a callable.

[zzzeek]

Right, it probably loads up the whole script environment first before it runs your env.py that actually sets up the contexts. still want this?

[zzzeek]

OK , I had a little more of an OO approach for this:

self._migration_context = MigrationContext.configure(
    connection=connection,
    url=url,
    dialect_name=dialect_name,
    opts=opts,
    environment_context=self
)

--- a/alembic/migration.py
+++ b/alembic/migration.py
@@ -110,6 +110,12 @@ class MigrationContext(object):
                         "transactional" if self.impl.transactional_ddl
                         else "non-transactional")

+    @property
+    def config(self):
+        """Return the :class:`.Config` used by the current environment, if any."""
+        return self.environment_context.config
+
+
     @classmethod
     def configure(cls,
                 connection=None,

etc. can we do it like that? thanks.

[msabramo]

Hmm, that's probably better. I will try it on the train in a bit.

[zzzeek]

in fact environment_context can be None so maybe do a check for that in the config property too.

[msabramo]

Just did an amend and push -f; see how that looks.

 + 3fcb2b2...4445c83 add_env_context_and_config_to_migration_context -> add_env_context_and_config_to_migration_context (forced update)

[msabramo]

Yeah, I worked around the problem by changing my migration to be lazy.

def get_schema():
    return op.get_context().config.get_main_option('schema')

def upgrade():
    schema = get_schema()
    print("upgrade (%r): schema = %r" % (revision, schema))
    op.execute("""
        INSERT INTO {schema}.sm_Bench_Source_D_Type
...

I do want something that allows me to access config from the migration so that I can insert this schema name. The exact mechanics could be different.

[msabramo]

I guess it would be nice if I didn't have to do the the substitution for .format(schema=schema)) in every command, but I don't know if there's any way to handle that in the tool, as you probably can't predict all the stuff people might want to do. Can I put my own execute function in env.py and then import it from the migrations?

    op.execute("""
        TRUNCATE TABLE {schema}.Source_Type;
        """.format(schema=schema))

Wondering if there could perhaps be some API by which I tell Alembic about some callback that transforms all the queries?

[zzzeek]

make environment_context optional here, so that existing usage of the MigrationContext by itself doesnt break

[msabramo]

Done.

[zzzeek]

looks great, it's on my list of PRs to pull in.

in the meantime, if you feel like adding a test, feel free :)

[zzzeek]

thanks, added a test + changelog in bd8a8dc.

[msabramo]

Awesome! Thank you! I was going to write a test, I swear :-)

What's the :pullreq: stuff in the Changelog. Intrigued...

[zzzeek]

see the link at http://alembic.readthedocs.org/en/latest/changelog.html#change-d057961ce5b3254ebbbc9fcf8e533680