From 9aa72c5cc29c45a79675da10c9a2d22e848f7cd8 Mon Sep 17 00:00:00 2001
From: Joseph Schorr <joseph.schorr@coreos.com>
Date: Mon, 3 Nov 2014 15:25:55 -0500
Subject: [PATCH] Fix migration issues:   - MySQL 5.5 doesn't support the now()
 call as a default   - Postgres migration isn't auto-committed, so we have to
 check if the table exists first

---
 data/migrations/env.py                           |  9 +++++++--
 ...8d_add_brute_force_prevention_metadata_to_.py |  2 +-
 util/uncompressedsize.py                         | 16 +++++++++++++---
 3 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/data/migrations/env.py b/data/migrations/env.py
index f27a483f7..650683572 100644
--- a/data/migrations/env.py
+++ b/data/migrations/env.py
@@ -5,7 +5,7 @@ import os
 from alembic import context
 from sqlalchemy import engine_from_config, pool
 from logging.config import fileConfig
-from urllib import unquote
+from urllib import unquote, quote
 from peewee import SqliteDatabase
 
 from data.database import all_models, db
@@ -24,6 +24,11 @@ if 'GENMIGRATE' in os.environ:
     else:
         db_uri = 'postgresql://postgres@%s/genschema' % (docker_host_ip)
 
+if 'DB_URI' in os.environ:
+    db_uri = os.environ['DB_URI']
+
+app.config['DB_URI'] = db_uri
+
 config = context.config
 config.set_main_option('sqlalchemy.url', db_uri)
 
@@ -69,7 +74,7 @@ def run_migrations_online():
 
     """
 
-    if isinstance(db.obj, SqliteDatabase) and not 'GENMIGRATE' in os.environ:
+    if isinstance(db.obj, SqliteDatabase) and not 'GENMIGRATE' in os.environ and not 'DB_URI' in os.environ:
         print ('Skipping Sqlite migration!')
         return
 
diff --git a/data/migrations/versions/4fdb65816b8d_add_brute_force_prevention_metadata_to_.py b/data/migrations/versions/4fdb65816b8d_add_brute_force_prevention_metadata_to_.py
index 1ce802eca..bc8373655 100644
--- a/data/migrations/versions/4fdb65816b8d_add_brute_force_prevention_metadata_to_.py
+++ b/data/migrations/versions/4fdb65816b8d_add_brute_force_prevention_metadata_to_.py
@@ -17,7 +17,7 @@ from sqlalchemy.dialects import mysql
 def upgrade(tables):
     ### commands auto generated by Alembic - please adjust! ###
     op.add_column('user', sa.Column('invalid_login_attempts', sa.Integer(), nullable=False, server_default="0"))
-    op.add_column('user', sa.Column('last_invalid_login', sa.DateTime(), nullable=False, server_default=sa.func.now()))
+    op.add_column('user', sa.Column('last_invalid_login', sa.DateTime(), nullable=False))
     ### end Alembic commands ###
 
 
diff --git a/util/uncompressedsize.py b/util/uncompressedsize.py
index bb12812b2..771a06abf 100644
--- a/util/uncompressedsize.py
+++ b/util/uncompressedsize.py
@@ -3,7 +3,7 @@ import zlib
 import sys
 
 from data import model
-from data.database import ImageStorage
+from data.database import ImageStorage, configure
 from app import app, storage as store
 from data.database import db, db_random_func
 from util.gzipstream import ZLIB_GZIP_WINDOW
@@ -14,11 +14,13 @@ logger = logging.getLogger(__name__)
 
 CHUNK_SIZE = 5 * 1024 * 1024
 
-
 def backfill_sizes_from_data():
   logger.setLevel(logging.DEBUG)
   logger.debug('Starting uncompressed image size backfill')
-  
+
+  # Make sure we have a reference to the current DB.
+  configure(app.config)
+
   formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
   
   ch = logging.StreamHandler(sys.stdout)
@@ -27,6 +29,14 @@ def backfill_sizes_from_data():
 
   encountered = set()
 
+  # Try reading the ImageStorage table count. If it doesn't exist, then this is a postgres
+  # initial setup migration and we can skip this step anyway.
+  try:
+    ImageStorage.select().count()
+  except:
+    logger.debug('Skipping migration for new setup')
+    return
+
   while True:
     # Load the record from the DB.
     batch_ids = list(ImageStorage