~ubuntuone-pqm-team/canonical-identity-provider/trunk

« back to all changes in this revision

Viewing changes to identityprovider/management/commands/sqlcachedloaddata.py

  • Committer: Danny Tamez
  • Date: 2010-04-21 15:29:24 UTC
  • Revision ID: danny.tamez@canonical.com-20100421152924-lq1m92tstk2iz75a
Canonical SSO Provider (Open Source) - Initial Commit

Show diffs side-by-side

added added

removed removed

Lines of Context:
 
1
# Copyright 2010 Canonical Ltd.  This software is licensed under the
 
2
# GNU Affero General Public License version 3 (see the file LICENSE).
 
3
 
 
4
# Parts of this module are based on code from the Django project.
 
5
# Please see the license file in the third-party/django directory
 
6
 
 
7
import os
 
8
import sys
 
9
 
 
10
from django.core.management.base import BaseCommand
 
11
from django.core.management.color import no_style
 
12
from django.utils.translation import ugettext as _
 
13
 
 
14
 
 
15
class Command(BaseCommand):
 
16
    fixture_cache = {}
 
17
    help = _('Installs the named fixture(s) in the database.')
 
18
    args = _("fixture [fixture ...]")
 
19
 
 
20
    def handle(self, *fixture_labels, **options):
 
21
        from django.db.models import get_apps
 
22
        from django.core import serializers
 
23
        from django.db import connection, transaction
 
24
        from django.conf import settings
 
25
        self.style = no_style()
 
26
 
 
27
        show_traceback = options.get('traceback', False)
 
28
 
 
29
        # commit is a stealth option - it isn't really useful as
 
30
        # a command line option, but it can be useful when invoking
 
31
        # loaddata from within another script.
 
32
        # If commit=True, loaddata will use its own transaction;
 
33
        # if commit=False, the data load SQL will become part of
 
34
        # the transaction in place when loaddata was invoked.
 
35
        commit = options.get('commit', True)
 
36
 
 
37
        # Keep a count of the installed objects and fixtures
 
38
        fixture_count = 0
 
39
        object_count = 0
 
40
        models = set()
 
41
 
 
42
        humanize = lambda dirname: dirname and "'%s'" % dirname or 'absolute path'
 
43
 
 
44
        # Get a cursor (even though we don't need one yet). This has
 
45
        # the side effect of initializing the test database (if
 
46
        # it isn't already initialized).
 
47
        cursor = connection.cursor()
 
48
        # Start transaction management. All fixtures are installed in a
 
49
        # single transaction to ensure that all references are resolved.
 
50
        if commit:
 
51
            transaction.commit_unless_managed()
 
52
            transaction.enter_transaction_management()
 
53
            transaction.managed(True)
 
54
 
 
55
        app_fixtures = [os.path.join(os.path.dirname(app.__file__), 'fixtures')
 
56
                        for app in get_apps()]
 
57
        for fixture_label in fixture_labels:
 
58
            if fixture_label in Command.fixture_cache:
 
59
                fixture = Command.fixture_cache[fixture_label]
 
60
                cursor.execute(fixture['sql'])
 
61
                object_count += 1
 
62
                models.update(fixture['models'])
 
63
                continue
 
64
            else:
 
65
                current_query = len(connection.queries)
 
66
 
 
67
            new_fixture = {'models': set()}
 
68
            parts = fixture_label.split('.')
 
69
            if len(parts) == 1:
 
70
                fixture_name = fixture_label
 
71
                formats = serializers.get_public_serializer_formats()
 
72
            else:
 
73
                fixture_name, format = '.'.join(parts[:-1]), parts[-1]
 
74
                if format in serializers.get_public_serializer_formats():
 
75
                    formats = [format]
 
76
                else:
 
77
                    formats = []
 
78
 
 
79
            if not formats:
 
80
                sys.stderr.write(
 
81
                    self.style.ERROR(_("Problem installing fixture "
 
82
                                       "'%(fixture_name)s': %(format)s is "
 
83
                                       "not a known serialization format.") %
 
84
                                     {'fixture_name': fixture_name,
 
85
                                      'format': format}))
 
86
                transaction.rollback()
 
87
                transaction.leave_transaction_management()
 
88
                return
 
89
 
 
90
            if os.path.isabs(fixture_name):
 
91
                fixture_dirs = [fixture_name]
 
92
            else:
 
93
                fixture_dirs = (app_fixtures +
 
94
                                list(settings.FIXTURE_DIRS) + [''])
 
95
 
 
96
            for fixture_dir in fixture_dirs:
 
97
                label_found = False
 
98
                for format in formats:
 
99
                    serializer = serializers.get_serializer(format)
 
100
                    try:
 
101
                        full_path = os.path.join(
 
102
                            fixture_dir, '.'.join([fixture_name, format]))
 
103
                        fixture = open(full_path, 'r')
 
104
                        if label_found:
 
105
                            fixture.close()
 
106
                            print self.style.ERROR(
 
107
                                "Multiple fixtures named '%(fixture_name)s'"
 
108
                                "in %(fixture_dir)s. Aborting." %
 
109
                                {'fixture_name': fixture_name,
 
110
                                 'fixture_dir': humanize(fixture_dir)})
 
111
                            transaction.rollback()
 
112
                            transaction.leave_transaction_management()
 
113
                            return
 
114
                        else:
 
115
                            fixture_count += 1
 
116
                            objects_in_fixture = 0
 
117
                            try:
 
118
                                objects = serializers.deserialize(format,
 
119
                                                                  fixture)
 
120
                                for obj in objects:
 
121
                                    objects_in_fixture += 1
 
122
                                    models.add(obj.object.__class__)
 
123
                                    new_fixture['models'].add(
 
124
                                        obj.object.__class__)
 
125
                                    obj.save()
 
126
                                object_count += objects_in_fixture
 
127
                                label_found = True
 
128
                            except (SystemExit, KeyboardInterrupt):
 
129
                                raise
 
130
                            except Exception:
 
131
                                import traceback
 
132
                                fixture.close()
 
133
                                transaction.rollback()
 
134
                                transaction.leave_transaction_management()
 
135
                                if show_traceback:
 
136
                                    traceback.print_exc()
 
137
                                else:
 
138
                                    sys.stderr.write(
 
139
                                        self.style.ERROR(
 
140
                                            "Problem installing fixture "
 
141
                                            "'%(full_path)s': %(exception)s\n" %
 
142
                                            {'full_path': full_path,
 
143
                                             'exception': ''.join(traceback.format_exception(
 
144
                                                 sys.exc_type,
 
145
                                                 sys.exc_value,
 
146
                                                 sys.exc_traceback))}))
 
147
                                return
 
148
                            fixture.close()
 
149
                            all = connection.queries[current_query:]
 
150
                            queries = [x['sql'] for x in all
 
151
                                       if not x['sql'].lower().startswith('select')]
 
152
                            new_fixture['sql'] = ';'.join(queries)
 
153
                            Command.fixture_cache[fixture_label] = new_fixture
 
154
 
 
155
                            # If the fixture we loaded contains 0 objects,
 
156
                            # assume that an error was encountered during
 
157
                            # fixture loading.
 
158
                            if objects_in_fixture == 0:
 
159
                                sys.stderr.write(
 
160
                                    self.style.ERROR(
 
161
                                        "No fixture data found for "
 
162
                                        "'%(fixture_name)s'. (File format "
 
163
                                        "may be invalid.)" %
 
164
                                        {'fixture_name': fixture_name}))
 
165
                                transaction.rollback()
 
166
                                transaction.leave_transaction_management()
 
167
                                return
 
168
                    except:
 
169
                        pass
 
170
 
 
171
        # If we found even one object in a fixture, we need to reset the
 
172
        # database sequences.
 
173
        if object_count > 0:
 
174
            sequence_sql = connection.ops.sequence_reset_sql(self.style,
 
175
                                                             models)
 
176
            if sequence_sql:
 
177
                for line in sequence_sql:
 
178
                    cursor.execute(line)
 
179
 
 
180
        if commit:
 
181
            transaction.commit()
 
182
            transaction.leave_transaction_management()
 
183
 
 
184
        # Close the DB connection. This is required as a workaround for an
 
185
        # edge case in MySQL: if the same connection is used to
 
186
        # create tables, load data, and query, the query can return
 
187
        # incorrect results. See Django #7572, MySQL #37735.
 
188
        if commit:
 
189
            connection.close()