From fbbb669b8cd8845cd5dbdd827f8beb70fa1476bc Mon Sep 17 00:00:00 2001 From: Mike Bayer Date: Thu, 6 Apr 2017 13:55:35 -0400 Subject: [PATCH] Add timezone option to config Using dateutil.tz to link string names to tzinfo objects, the create_date can now generate using a named timezone rather than datetime.now(). Change-Id: I9f151cb9e11da3d68be63d7141f60e7eccb9812c Fixes: #425 --- alembic/script/base.py | 23 +++++++- alembic/templates/generic/alembic.ini.mako | 6 ++ alembic/templates/multidb/alembic.ini.mako | 6 ++ alembic/templates/pylons/alembic.ini.mako | 6 ++ docs/build/changelog.rst | 10 ++++ docs/build/tutorial.rst | 25 +++++++- setup.py | 1 + tests/test_script_production.py | 67 ++++++++++++++++++++++ 8 files changed, 139 insertions(+), 5 deletions(-) diff --git a/alembic/script/base.py b/alembic/script/base.py index a79ec09..17cb3de 100644 --- a/alembic/script/base.py +++ b/alembic/script/base.py @@ -1,4 +1,5 @@ import datetime +from dateutil import tz import os import re import shutil @@ -42,7 +43,8 @@ class ScriptDirectory(object): def __init__(self, dir, file_template=_default_file_template, truncate_slug_length=40, version_locations=None, - sourceless=False, output_encoding="utf-8"): + sourceless=False, output_encoding="utf-8", + timezone=None): self.dir = dir self.file_template = file_template self.version_locations = version_locations @@ -50,6 +52,7 @@ class ScriptDirectory(object): self.sourceless = sourceless self.output_encoding = output_encoding self.revision_map = revision.RevisionMap(self._load_revisions) + self.timezone = timezone if not os.access(dir, os.F_OK): raise util.CommandError("Path doesn't exist: %r. Please use " @@ -118,6 +121,7 @@ class ScriptDirectory(object): version_locations = config.get_main_option("version_locations") if version_locations: version_locations = _split_on_space_comma.split(version_locations) + return ScriptDirectory( util.coerce_resource_to_filename(script_location), file_template=config.get_main_option( @@ -126,7 +130,8 @@ class ScriptDirectory(object): truncate_slug_length=truncate_slug_length, sourceless=config.get_main_option("sourceless") == "true", output_encoding=config.get_main_option("output_encoding", "utf-8"), - version_locations=version_locations + version_locations=version_locations, + timezone=config.get_main_option("timezone") ) @contextmanager @@ -440,6 +445,18 @@ class ScriptDirectory(object): "Creating directory %s" % path, os.makedirs, path) + def _generate_create_date(self): + if self.timezone is not None: + tzinfo = tz.gettz(self.timezone.upper()) + if tzinfo is None: + raise util.CommandError( + "Can't locate timezone: %s" % self.timezone) + create_date = datetime.datetime.utcnow().replace( + tzinfo=tz.tzutc()).astimezone(tzinfo) + else: + create_date = datetime.datetime.now() + return create_date + def generate_revision( self, revid, message, head=None, refresh=False, splice=False, branch_labels=None, @@ -478,7 +495,7 @@ class ScriptDirectory(object): if len(set(heads)) != len(heads): raise util.CommandError("Duplicate head revisions specified") - create_date = datetime.datetime.now() + create_date = self._generate_create_date() if version_path is None: if len(self._version_locations) > 1: diff --git a/alembic/templates/generic/alembic.ini.mako b/alembic/templates/generic/alembic.ini.mako index 4d3bf6e..9ee59db 100644 --- a/alembic/templates/generic/alembic.ini.mako +++ b/alembic/templates/generic/alembic.ini.mako @@ -7,6 +7,12 @@ script_location = ${script_location} # template used to generate migration files # file_template = %%(rev)s_%%(slug)s +# timezone to use when rendering the date +# within the migration file as well as the filename. +# string value is passed to dateutil.tz.gettz() +# leave blank for localtime +# timezone = + # max length of characters to apply to the # "slug" field #truncate_slug_length = 40 diff --git a/alembic/templates/multidb/alembic.ini.mako b/alembic/templates/multidb/alembic.ini.mako index 929a4be..a0708ff 100644 --- a/alembic/templates/multidb/alembic.ini.mako +++ b/alembic/templates/multidb/alembic.ini.mako @@ -7,6 +7,12 @@ script_location = ${script_location} # template used to generate migration files # file_template = %%(rev)s_%%(slug)s +# timezone to use when rendering the date +# within the migration file as well as the filename. +# string value is passed to dateutil.tz.gettz() +# leave blank for localtime +# timezone = + # max length of characters to apply to the # "slug" field #truncate_slug_length = 40 diff --git a/alembic/templates/pylons/alembic.ini.mako b/alembic/templates/pylons/alembic.ini.mako index 62191e0..c5cc413 100644 --- a/alembic/templates/pylons/alembic.ini.mako +++ b/alembic/templates/pylons/alembic.ini.mako @@ -7,6 +7,12 @@ script_location = ${script_location} # template used to generate migration files # file_template = %%(rev)s_%%(slug)s +# timezone to use when rendering the date +# within the migration file as well as the filename. +# string value is passed to dateutil.tz.gettz() +# leave blank for localtime +# timezone = + # max length of characters to apply to the # "slug" field #truncate_slug_length = 40 diff --git a/docs/build/changelog.rst b/docs/build/changelog.rst index 7835883..926556b 100644 --- a/docs/build/changelog.rst +++ b/docs/build/changelog.rst @@ -7,6 +7,16 @@ Changelog :version: 0.9.2 :released: + .. change:: 425 + :tags: feature, commands + :tickets: 425 + + Added a new configuration option ``timezone``, a string timezone name + that will be applied to the create date timestamp rendered + inside the revision file as made availble to the ``file_template`` used + to generate the revision filename. Note this change adds the + ``python-dateutil`` package as a dependency. + .. change:: 421 :tags: bug, autogenerate :tickets: 421 diff --git a/docs/build/tutorial.rst b/docs/build/tutorial.rst index fbbc5a4..2f71b02 100644 --- a/docs/build/tutorial.rst +++ b/docs/build/tutorial.rst @@ -119,6 +119,12 @@ The file generated with the "generic" configuration looks like:: # template used to generate migration files # file_template = %%(rev)s_%%(slug)s + # timezone to use when rendering the date + # within the migration file as well as the filename. + # string value is passed to dateutil.tz.gettz() + # leave blank for localtime + # timezone = + # max length of characters to apply to the # "slug" field #truncate_slug_length = 40 @@ -212,8 +218,23 @@ This file contains the following features: * ``%%(rev)s`` - revision id * ``%%(slug)s`` - a truncated string derived from the revision message * ``%%(year)d``, ``%%(month).2d``, ``%%(day).2d``, ``%%(hour).2d``, - ``%%(minute).2d``, ``%%(second).2d`` - components of the create date - as returned by ``datetime.datetime.now()`` + ``%%(minute).2d``, ``%%(second).2d`` - components of the create date, + by default ``datetime.datetime.now()`` unless the ``timezone`` + configuration option is also used. + +* ``timezone`` - an optional timezone name (e.g. ``UTC``, ``EST5EDT``, etc.) + that will be applied to the timestamp which renders inside the migration + file's comment as well as within the filename. If ``timezone`` is specified, + the create date object is no longer derived from ``datetime.datetime.now()`` + and is instead generated as:: + + datetime.datetime.utcnow().replace( + tzinfo=dateutil.tz.tzutc() + ).astimezone( + dateutil.tz.gettz() + ) + + .. versionadded:: 0.9.2 * ``truncate_slug_length`` - defaults to 40, the max number of characters to include in the "slug" field. diff --git a/setup.py b/setup.py index 445e335..e424eb0 100644 --- a/setup.py +++ b/setup.py @@ -16,6 +16,7 @@ requires = [ 'SQLAlchemy>=0.7.6', 'Mako', 'python-editor>=0.3', + 'python-dateutil' ] try: diff --git a/tests/test_script_production.py b/tests/test_script_production.py index 3703014..3364d56 100644 --- a/tests/test_script_production.py +++ b/tests/test_script_production.py @@ -18,6 +18,7 @@ import sqlalchemy as sa from sqlalchemy.engine.reflection import Inspector from alembic.util import CommandError import re +from dateutil import tz env, abc, def_ = None, None, None @@ -153,6 +154,72 @@ class ScriptNamingTest(TestBase): "message_2012_7_25_15_8_5.py" % _get_staging_directory()) ) + def _test_tz(self, timezone_arg, given, expected): + script = ScriptDirectory( + _get_staging_directory(), + file_template="%(rev)s_%(slug)s_" + "%(year)s_%(month)s_" + "%(day)s_%(hour)s_" + "%(minute)s_%(second)s", + timezone=timezone_arg + ) + + with mock.patch( + "alembic.script.base.datetime", + mock.Mock( + datetime=mock.Mock( + utcnow=lambda: given, + now=lambda: given + ) + ) + ): + create_date = script._generate_create_date() + eq_( + create_date, + expected + ) + + def test_custom_tz(self): + self._test_tz( + 'EST5EDT', + datetime.datetime(2012, 7, 25, 15, 8, 5), + datetime.datetime( + 2012, 7, 25, 11, 8, 5, tzinfo=tz.gettz('EST5EDT')) + ) + + def test_custom_tz_lowercase(self): + self._test_tz( + 'est5edt', + datetime.datetime(2012, 7, 25, 15, 8, 5), + datetime.datetime( + 2012, 7, 25, 11, 8, 5, tzinfo=tz.gettz('EST5EDT')) + ) + + def test_custom_tz_utc(self): + self._test_tz( + 'utc', + datetime.datetime(2012, 7, 25, 15, 8, 5), + datetime.datetime( + 2012, 7, 25, 15, 8, 5, tzinfo=tz.gettz('UTC')) + ) + + def test_default_tz(self): + self._test_tz( + None, + datetime.datetime(2012, 7, 25, 15, 8, 5), + datetime.datetime(2012, 7, 25, 15, 8, 5) + ) + + def test_tz_cant_locate(self): + assert_raises_message( + CommandError, + "Can't locate timezone: fake", + self._test_tz, + "fake", + datetime.datetime(2012, 7, 25, 15, 8, 5), + datetime.datetime(2012, 7, 25, 15, 8, 5) + ) + class RevisionCommandTest(TestBase): def setUp(self):