From e6240c56fc907eba960cc9a08d7a15aee8fa9ff5 Mon Sep 17 00:00:00 2001
From: Nimisha Asthagiri <nasthagiri@edx.org>
Date: Thu, 18 Feb 2016 18:33:31 -0500
Subject: [PATCH] Update generate_course_blocks management command.

---
 .../commands/generate_course_blocks.py        | 100 +++++++++++++++++-
 1 file changed, 95 insertions(+), 5 deletions(-)

diff --git a/lms/djangoapps/course_blocks/management/commands/generate_course_blocks.py b/lms/djangoapps/course_blocks/management/commands/generate_course_blocks.py
index 6f3bfa3c62..2967b01316 100644
--- a/lms/djangoapps/course_blocks/management/commands/generate_course_blocks.py
+++ b/lms/djangoapps/course_blocks/management/commands/generate_course_blocks.py
@@ -1,6 +1,7 @@
 """
 Command to load course blocks.
 """
+from collections import defaultdict
 import logging
 
 from django.core.management.base import BaseCommand, CommandError
@@ -45,11 +46,32 @@ class Command(BaseCommand):
             action='store_true',
             default=False,
         )
+        parser.add_argument(
+            '--verbose',
+            help='Enable verbose logging.',
+            action='store_true',
+            default=False,
+        )
+        parser.add_argument(
+            '--start',
+            help='Starting index of course.',
+            default=0,
+            type=int,
+        )
+        parser.add_argument(
+            '--end',
+            help='Ending index of course.',
+            default=0,
+            type=int,
+        )
 
     def handle(self, *args, **options):
 
         if options.get('all'):
             course_keys = [course.id for course in modulestore().get_course_summaries()]
+            if options.get('start'):
+                end = options.get('end') or len(course_keys)
+                course_keys = course_keys[options['start']:end]
         else:
             if len(args) < 1:
                 raise CommandError('At least one course or --all must be specified.')
@@ -59,8 +81,13 @@ class Command(BaseCommand):
                 raise CommandError('Invalid key specified.')
 
         log.info('Generating course blocks for %d courses.', len(course_keys))
-        log.debug('Generating course blocks for the following courses: %s', course_keys)
 
+        if options.get('verbose'):
+            log.setLevel(logging.DEBUG)
+        else:
+            log.setLevel(logging.CRITICAL)
+
+        dag_info = _DAGInfo()
         for course_key in course_keys:
             try:
                 if options.get('force'):
@@ -68,7 +95,7 @@ class Command(BaseCommand):
                 else:
                     block_structure = get_course_in_cache(course_key)
                 if options.get('dags'):
-                    self._find_and_log_dags(block_structure, course_key)
+                    self._find_and_log_dags(block_structure, course_key, dag_info)
             except Exception as ex:  # pylint: disable=broad-except
                 log.exception(
                     'An error occurred while generating course blocks for %s: %s',
@@ -78,20 +105,83 @@ class Command(BaseCommand):
 
         log.info('Finished generating course blocks.')
 
-    def _find_and_log_dags(self, block_structure, course_key):
+        if options.get('dags'):
+            log.critical('DAG data: %s', unicode(dag_info))
+
+    def _find_and_log_dags(self, block_structure, course_key, dag_info):
         """
         Finds all DAGs within the given block structure.
 
         Arguments:
             BlockStructureBlockData - The block structure in which to find DAGs.
         """
-        log.info('DAG check starting for course %s.', unicode(course_key))
         for block_key in block_structure.get_block_keys():
             parents = block_structure.get_parents(block_key)
             if len(parents) > 1:
+                dag_info.on_dag_found(course_key, block_key)
                 log.warning(
                     'DAG alert - %s has multiple parents: %s.',
                     unicode(block_key),
                     [unicode(parent) for parent in parents],
                 )
-        log.info('DAG check complete for course %s.', unicode(course_key))
+
+
+class PrettyDefaultDict(defaultdict):
+    """
+    Wraps defaultdict to provide a better string representation.
+    """
+    __repr__ = dict.__repr__
+
+
+class _DAGBlockTypeInfo(object):
+    """
+    Class for aggregated DAG data for a specific block type.
+    """
+    def __init__(self):
+        self.num_of_dag_blocks = 0
+
+    def __repr__(self):
+        return repr(vars(self))
+
+
+class _DAGCourseInfo(object):
+    """
+    Class for aggregated DAG data for a specific course run.
+    """
+    def __init__(self):
+        self.num_of_dag_blocks = 0
+        self.dag_data_by_block_type = PrettyDefaultDict(_DAGBlockTypeInfo)
+
+    def __repr__(self):
+        return repr(vars(self))
+
+    def on_dag_found(self, block_key):
+        """
+        Updates DAG collected data for the given block.
+        """
+        self.num_of_dag_blocks += 1
+        self.dag_data_by_block_type[block_key.category].num_of_dag_blocks += 1
+
+
+class _DAGInfo(object):
+    """
+    Class for aggregated DAG data.
+    """
+    def __init__(self):
+        self.total_num_of_dag_blocks = 0
+        self.total_num_of_dag_courses = 0
+        self.dag_data_by_course = PrettyDefaultDict(_DAGCourseInfo)
+        self.dag_data_by_block_type = PrettyDefaultDict(_DAGBlockTypeInfo)
+
+    def __repr__(self):
+        return repr(vars(self))
+
+    def on_dag_found(self, course_key, block_key):
+        """
+        Updates DAG collected data for the given block.
+        """
+        self.total_num_of_dag_blocks += 1
+        if course_key not in self.dag_data_by_course:
+            self.total_num_of_dag_courses += 1
+        self.dag_data_by_course[unicode(course_key)].on_dag_found(block_key)
+        self.dag_data_by_block_type[block_key.category].num_of_dag_blocks += 1