From 6eaf1dbb3fcac227c755ba45ee3a821c0afd12b0 Mon Sep 17 00:00:00 2001
From: Joseph Schorr <joseph.schorr@coreos.com>
Date: Wed, 22 Apr 2015 17:11:08 -0400
Subject: [PATCH] Make the repositoryactioncount worker disconnect from the DB
 between runs

---
 workers/repositoryactioncounter.py | 58 +++++++++++++++---------------
 1 file changed, 30 insertions(+), 28 deletions(-)

diff --git a/workers/repositoryactioncounter.py b/workers/repositoryactioncounter.py
index 5a73c194b..eed5f3420 100644
--- a/workers/repositoryactioncounter.py
+++ b/workers/repositoryactioncounter.py
@@ -3,7 +3,8 @@ import logging
 from apscheduler.schedulers.blocking import BlockingScheduler
 
 from app import app
-from data.database import Repository, LogEntry, RepositoryActionCount, db_random_func, fn
+from data.database import (Repository, LogEntry, RepositoryActionCount, db_random_func, fn,
+                           UseThenDisconnect)
 from datetime import date, datetime, timedelta
 
 POLL_PERIOD_SECONDS = 30
@@ -15,36 +16,37 @@ sched = BlockingScheduler()
 def count_repository_actions():
   """ Counts actions for a random repository for the previous day. """
 
-  try:
-    # Get a random repository to count.
-    today = date.today()
-    yesterday = today - timedelta(days=1)
-    has_yesterday_actions = (RepositoryActionCount.select(RepositoryActionCount.repository)
-                                                  .where(RepositoryActionCount.date == yesterday))
-
-    to_count = (Repository.select()
-                          .where(~(Repository.id << (has_yesterday_actions)))
-                          .order_by(db_random_func()).get())
-
-    logger.debug('Counting: %s', to_count.id)
-
-    actions = (LogEntry.select()
-                       .where(LogEntry.repository == to_count,
-                              LogEntry.datetime >= yesterday,
-                              LogEntry.datetime < today)
-                       .count())
-
-    # Create the row.
+  with UseThenDisconnect(app.config):
     try:
-      RepositoryActionCount.create(repository=to_count, date=yesterday, count=actions)
-    except:
-      logger.exception('Exception when writing count')
+      # Get a random repository to count.
+      today = date.today()
+      yesterday = today - timedelta(days=1)
+      has_yesterday_actions = (RepositoryActionCount.select(RepositoryActionCount.repository)
+                                                    .where(RepositoryActionCount.date == yesterday))
 
-    return True
+      to_count = (Repository.select()
+                            .where(~(Repository.id << (has_yesterday_actions)))
+                            .order_by(db_random_func()).get())
 
-  except Repository.DoesNotExist:
-    logger.debug('No further repositories to count')
-    return False
+      logger.debug('Counting: %s', to_count.id)
+
+      actions = (LogEntry.select()
+                         .where(LogEntry.repository == to_count,
+                                LogEntry.datetime >= yesterday,
+                                LogEntry.datetime < today)
+                         .count())
+
+      # Create the row.
+      try:
+        RepositoryActionCount.create(repository=to_count, date=yesterday, count=actions)
+      except:
+        logger.exception('Exception when writing count')
+
+      return True
+
+    except Repository.DoesNotExist:
+      logger.debug('No further repositories to count')
+      return False
 
 
 if __name__ == "__main__":