from collections import defaultdict

from django.core.management.base import BaseCommand, CommandError
from django.contrib.sessions.models import Session
from django.contrib.sessions.backends.base import SessionBase

class Command(BaseCommand):
    args = '--batch-size=n --bigger-than=n'
    help = 'Analyze Django sessions, summarizing present keys & session size'

    # XXX: Most of these should be set at __init__, not as class-level
    # constants
    # Number of sessions to process in one batch.
    # Tune as necessary to get DB load at a sweet spot.
    batch_size = 5000

    # Only process sessions with more than this many bytes.
    bigger_than = 10 * 1024

    # Total # of matching sessions
    processed_session_count = 0

    # Total # of sessions in DB
    total_session_count = 0

    # Encoded data sizes for the sessions we looked at
    sizes = []

    # Observed keys (key => frequency)
    keys = defaultdict(int)

    # Sizes of observed keys
    key_sizes = defaultdict(list)

    def handle(self, *args, **options):
        # process options.
        self.process_options(options)

        # Count all sessions
        self.total_session_count = Session.objects.all().count()

        # process sessions
        for session in self.get_sessions():
            self.process_session(session)

        # summarize results and exit
        self.print_results()

    def process_options(self, options):
        """
        Validate options & configure self appropriately based on option
        settings
        """
        if 'batch-size' in options:
            try:
                new_batch_size = int(options['batch-size'])
            except ValueError, e:
                raise CommandError('batch-size must be a positive integer')

            self.batch_size = new_batch_size

            if self.batch_size <= 0:
                raise CommandError('batch-size must be a positive integer')

        if 'bigger-than' in options:
            try:
                new_bigger_than = int(options['bigger-than'])
            except ValueError, e:
                raise CommandError('bigger-than must be an integer')

            self.bigger_than = new_bigger_than

    def get_filtered_queryset(self):
        """
        Return a Session qs with any configured filters already applied.
        """
        # XXX: Potentially MySQL-specific. Should be revised to handle
        #      other DBMS, or fail gracefully when they're in use.
        return Session.objects.extra(where=['LENGTH(session_data) > %d' % self.bigger_than])

    def get_sessions(self):
        """
        Return sessions that match the configured bigger than parameters.

        Internally, collect sessions in batches of size batch_size, then yield
        to the caller.
        """
        start = 0
        while True:
            qs = self.get_filtered_queryset()
            if qs.count() == 0 or start >= qs.count():
                return

            qs = qs[start:start+self.batch_size]
            start += self.batch_size

            for session in qs:
                yield session

    def process_session(self, session):
        """
        Process a session.
        """
        self.processed_session_count += 1

        data = session.session_data
        self.sizes.append(len(data))

        decoded = session.get_decoded()
        for key, value in decoded.iteritems():
            self.keys[key] += 1
            self.key_sizes[key].append(self.get_size(key, value))

    def get_size(self, key, value):
        """
        Serialize an object hierarchy & return the size of the result.
        """
        d = {key: value}
        d_enc = SessionBase().encode(d)
        return len(d_enc)

    def print_results(self):
        self.stdout.write("Processed %d sessions out of %d "
                          "total sessions\n" %
                            (self.processed_session_count,
                             self.total_session_count))

        average = sum(self.sizes) / float(len(self.sizes))
        self.stdout.write("Average size was %f bytes\n" % average)

        self.stdout.write("Saw the following keys:\n")
        for key, count in self.keys.iteritems():
            avg_size = sum(self.key_sizes[key]) / float(count)
            self.stdout.write("    %s (%d times, avg. size %f)\n"
                % (key, count, avg_size))

@@ -28,16 +28,49 @@
 aultdict
+%0Afrom optparse import make_option
 %0A%0Afrom d
@@ -402,308 +402,586 @@
     
-# XXX: Most of these should be set at __init__, not as class-level%0A    # constants%0A    # Number of sessions to process in one batch.%0A    # Tune as necessary to get DB load at a sweet spot.%0A    batch_size = 5000%0A%0A    # Only process sessions with more than this many bytes.%0A    bigger_than = 10 * 1024%0A
+option_list = BaseCommand.option_list + (%0A        make_option(%22--bigger-than%22,%0A            action='store',%0A            dest='bigger_than',%0A            type='int',%0A            default=10*1024,%0A            help=%22Only return records with more than BIGGER_THAN bytes%22),%0A        make_option(%22--batch-size%22,%0A            action='store',%0A            dest='batch_size',%0A            type='int',%0A            default=5000,%0A            help=%22Process sessions in batches no larger than BATCH_SIZE%22)%0A    )%0A%0A    # XXX: Most of these should be set at __init__, not as class-level%0A    # constants
 %0A   
@@ -1856,73 +1856,13 @@
     
-if 'batch-size' in options:%0A            try:%0A                new_
+self.
 batc
@@ -1870,20 +1870,16 @@
 _size = 
-int(
 options%5B
@@ -1888,16 +1888,15 @@
 atch
--
+_
 size'%5D
-)
 %0A   
@@ -1904,345 +1904,13 @@
     
-    except ValueError, e:%0A                raise CommandError('batch-size must be a positive integer')%0A%0A            self.batch_size = new_batch_size%0A%0A            if self.batch_size %3C= 0:%0A                raise CommandError('batch-size must be a positive integer')%0A%0A        if 'bigger-than' in options:%0A            try:%0A                new_
+self.
 bigg
@@ -1923,12 +1923,8 @@
 n = 
-int(
 opti
@@ -1938,167 +1938,15 @@
 gger
--
+_
 than'%5D
-)%0A            except ValueError, e:%0A                raise CommandError('bigger-than must be an integer')%0A%0A            self.bigger_than = new_bigger_than
 %0A%0A  

Fix broken option parsing.