Ticket #14001: creation.py

File creation.py, 20.1 KB (added by mnbayazit, 6 years ago)

hack for postgresql to allow choosing test database/username/password

Line 
1import sys
2import time
3
4from django.conf import settings
5from django.core.management import call_command
6
7# The prefix to put on the default database name when creating
8# the test database.
9TEST_DATABASE_PREFIX = 'test_'
10
11class BaseDatabaseCreation(object):
12    """
13    This class encapsulates all backend-specific differences that pertain to
14    database *creation*, such as the column types to use for particular Django
15    Fields, the SQL used to create and destroy tables, and the creation and
16    destruction of test databases.
17    """
18    data_types = {}
19
20    def __init__(self, connection):
21        self.connection = connection
22
23    def _digest(self, *args):
24        """
25        Generates a 32-bit digest of a set of arguments that can be used to
26        shorten identifying names.
27        """
28        return '%x' % (abs(hash(args)) % 4294967296L)  # 2**32
29
30    def sql_create_model(self, model, style, known_models=set()):
31        """
32        Returns the SQL required to create a single model, as a tuple of:
33            (list_of_sql, pending_references_dict)
34        """
35        from django.db import models
36
37        opts = model._meta
38        if not opts.managed or opts.proxy:
39            return [], {}
40        final_output = []
41        table_output = []
42        pending_references = {}
43        qn = self.connection.ops.quote_name
44        for f in opts.local_fields:
45            col_type = f.db_type(connection=self.connection)
46            tablespace = f.db_tablespace or opts.db_tablespace
47            if col_type is None:
48                # Skip ManyToManyFields, because they're not represented as
49                # database columns in this table.
50                continue
51            # Make the definition (e.g. 'foo VARCHAR(30)') for this field.
52            field_output = [style.SQL_FIELD(qn(f.column)),
53                style.SQL_COLTYPE(col_type)]
54            if not f.null:
55                field_output.append(style.SQL_KEYWORD('NOT NULL'))
56            if f.primary_key:
57                field_output.append(style.SQL_KEYWORD('PRIMARY KEY'))
58            elif f.unique:
59                field_output.append(style.SQL_KEYWORD('UNIQUE'))
60            if tablespace and f.unique:
61                # We must specify the index tablespace inline, because we
62                # won't be generating a CREATE INDEX statement for this field.
63                field_output.append(self.connection.ops.tablespace_sql(tablespace, inline=True))
64            if f.rel:
65                ref_output, pending = self.sql_for_inline_foreign_key_references(f, known_models, style)
66                if pending:
67                    pr = pending_references.setdefault(f.rel.to, []).append((model, f))
68                else:
69                    field_output.extend(ref_output)
70            table_output.append(' '.join(field_output))
71        for field_constraints in opts.unique_together:
72            table_output.append(style.SQL_KEYWORD('UNIQUE') + ' (%s)' % \
73                ", ".join([style.SQL_FIELD(qn(opts.get_field(f).column)) for f in field_constraints]))
74
75        full_statement = [style.SQL_KEYWORD('CREATE TABLE') + ' ' + style.SQL_TABLE(qn(opts.db_table)) + ' (']
76        for i, line in enumerate(table_output): # Combine and add commas.
77            full_statement.append('    %s%s' % (line, i < len(table_output)-1 and ',' or ''))
78        full_statement.append(')')
79        if opts.db_tablespace:
80            full_statement.append(self.connection.ops.tablespace_sql(opts.db_tablespace))
81        full_statement.append(';')
82        final_output.append('\n'.join(full_statement))
83
84        if opts.has_auto_field:
85            # Add any extra SQL needed to support auto-incrementing primary keys.
86            auto_column = opts.auto_field.db_column or opts.auto_field.name
87            autoinc_sql = self.connection.ops.autoinc_sql(opts.db_table, auto_column)
88            if autoinc_sql:
89                for stmt in autoinc_sql:
90                    final_output.append(stmt)
91
92        return final_output, pending_references
93
94    def sql_for_inline_foreign_key_references(self, field, known_models, style):
95        "Return the SQL snippet defining the foreign key reference for a field"
96        qn = self.connection.ops.quote_name
97        if field.rel.to in known_models:
98            output = [style.SQL_KEYWORD('REFERENCES') + ' ' + \
99                style.SQL_TABLE(qn(field.rel.to._meta.db_table)) + ' (' + \
100                style.SQL_FIELD(qn(field.rel.to._meta.get_field(field.rel.field_name).column)) + ')' +
101                self.connection.ops.deferrable_sql()
102            ]
103            pending = False
104        else:
105            # We haven't yet created the table to which this field
106            # is related, so save it for later.
107            output = []
108            pending = True
109
110        return output, pending
111
112    def sql_for_pending_references(self, model, style, pending_references):
113        "Returns any ALTER TABLE statements to add constraints after the fact."
114        from django.db.backends.util import truncate_name
115
116        if not model._meta.managed or model._meta.proxy:
117            return []
118        qn = self.connection.ops.quote_name
119        final_output = []
120        opts = model._meta
121        if model in pending_references:
122            for rel_class, f in pending_references[model]:
123                rel_opts = rel_class._meta
124                r_table = rel_opts.db_table
125                r_col = f.column
126                table = opts.db_table
127                col = opts.get_field(f.rel.field_name).column
128                # For MySQL, r_name must be unique in the first 64 characters.
129                # So we are careful with character usage here.
130                r_name = '%s_refs_%s_%s' % (r_col, col, self._digest(r_table, table))
131                final_output.append(style.SQL_KEYWORD('ALTER TABLE') + ' %s ADD CONSTRAINT %s FOREIGN KEY (%s) REFERENCES %s (%s)%s;' % \
132                    (qn(r_table), qn(truncate_name(r_name, self.connection.ops.max_name_length())),
133                    qn(r_col), qn(table), qn(col),
134                    self.connection.ops.deferrable_sql()))
135            del pending_references[model]
136        return final_output
137
138    def sql_for_many_to_many(self, model, style):
139        "Return the CREATE TABLE statments for all the many-to-many tables defined on a model"
140        import warnings
141        warnings.warn(
142            'Database creation API for m2m tables has been deprecated. M2M models are now automatically generated',
143            PendingDeprecationWarning
144        )
145
146        output = []
147        for f in model._meta.local_many_to_many:
148            if model._meta.managed or f.rel.to._meta.managed:
149                output.extend(self.sql_for_many_to_many_field(model, f, style))
150        return output
151
152    def sql_for_many_to_many_field(self, model, f, style):
153        "Return the CREATE TABLE statements for a single m2m field"
154        import warnings
155        warnings.warn(
156            'Database creation API for m2m tables has been deprecated. M2M models are now automatically generated',
157            PendingDeprecationWarning
158        )
159
160        from django.db import models
161        from django.db.backends.util import truncate_name
162
163        output = []
164        if f.auto_created:
165            opts = model._meta
166            qn = self.connection.ops.quote_name
167            tablespace = f.db_tablespace or opts.db_tablespace
168            if tablespace:
169                sql = self.connection.ops.tablespace_sql(tablespace, inline=True)
170                if sql:
171                    tablespace_sql = ' ' + sql
172                else:
173                    tablespace_sql = ''
174            else:
175                tablespace_sql = ''
176            table_output = [style.SQL_KEYWORD('CREATE TABLE') + ' ' + \
177                style.SQL_TABLE(qn(f.m2m_db_table())) + ' (']
178            table_output.append('    %s %s %s%s,' %
179                (style.SQL_FIELD(qn('id')),
180                style.SQL_COLTYPE(models.AutoField(primary_key=True).db_type(connection=self.connection)),
181                style.SQL_KEYWORD('NOT NULL PRIMARY KEY'),
182                tablespace_sql))
183
184            deferred = []
185            inline_output, deferred = self.sql_for_inline_many_to_many_references(model, f, style)
186            table_output.extend(inline_output)
187
188            table_output.append('    %s (%s, %s)%s' %
189                (style.SQL_KEYWORD('UNIQUE'),
190                style.SQL_FIELD(qn(f.m2m_column_name())),
191                style.SQL_FIELD(qn(f.m2m_reverse_name())),
192                tablespace_sql))
193            table_output.append(')')
194            if opts.db_tablespace:
195                # f.db_tablespace is only for indices, so ignore its value here.
196                table_output.append(self.connection.ops.tablespace_sql(opts.db_tablespace))
197            table_output.append(';')
198            output.append('\n'.join(table_output))
199
200            for r_table, r_col, table, col in deferred:
201                r_name = '%s_refs_%s_%s' % (r_col, col, self._digest(r_table, table))
202                output.append(style.SQL_KEYWORD('ALTER TABLE') + ' %s ADD CONSTRAINT %s FOREIGN KEY (%s) REFERENCES %s (%s)%s;' %
203                (qn(r_table),
204                qn(truncate_name(r_name, self.connection.ops.max_name_length())),
205                qn(r_col), qn(table), qn(col),
206                self.connection.ops.deferrable_sql()))
207
208            # Add any extra SQL needed to support auto-incrementing PKs
209            autoinc_sql = self.connection.ops.autoinc_sql(f.m2m_db_table(), 'id')
210            if autoinc_sql:
211                for stmt in autoinc_sql:
212                    output.append(stmt)
213        return output
214
215    def sql_for_inline_many_to_many_references(self, model, field, style):
216        "Create the references to other tables required by a many-to-many table"
217        import warnings
218        warnings.warn(
219            'Database creation API for m2m tables has been deprecated. M2M models are now automatically generated',
220            PendingDeprecationWarning
221        )
222
223        from django.db import models
224        opts = model._meta
225        qn = self.connection.ops.quote_name
226
227        table_output = [
228            '    %s %s %s %s (%s)%s,' %
229                (style.SQL_FIELD(qn(field.m2m_column_name())),
230                style.SQL_COLTYPE(models.ForeignKey(model).db_type(connection=self.connection)),
231                style.SQL_KEYWORD('NOT NULL REFERENCES'),
232                style.SQL_TABLE(qn(opts.db_table)),
233                style.SQL_FIELD(qn(opts.pk.column)),
234                self.connection.ops.deferrable_sql()),
235            '    %s %s %s %s (%s)%s,' %
236                (style.SQL_FIELD(qn(field.m2m_reverse_name())),
237                style.SQL_COLTYPE(models.ForeignKey(field.rel.to).db_type(connection=self.connection)),
238                style.SQL_KEYWORD('NOT NULL REFERENCES'),
239                style.SQL_TABLE(qn(field.rel.to._meta.db_table)),
240                style.SQL_FIELD(qn(field.rel.to._meta.pk.column)),
241                self.connection.ops.deferrable_sql())
242        ]
243        deferred = []
244
245        return table_output, deferred
246
247    def sql_indexes_for_model(self, model, style):
248        "Returns the CREATE INDEX SQL statements for a single model"
249        if not model._meta.managed or model._meta.proxy:
250            return []
251        output = []
252        for f in model._meta.local_fields:
253            output.extend(self.sql_indexes_for_field(model, f, style))
254        return output
255
256    def sql_indexes_for_field(self, model, f, style):
257        "Return the CREATE INDEX SQL statements for a single model field"
258        from django.db.backends.util import truncate_name
259
260        if f.db_index and not f.unique:
261            qn = self.connection.ops.quote_name
262            tablespace = f.db_tablespace or model._meta.db_tablespace
263            if tablespace:
264                sql = self.connection.ops.tablespace_sql(tablespace)
265                if sql:
266                    tablespace_sql = ' ' + sql
267                else:
268                    tablespace_sql = ''
269            else:
270                tablespace_sql = ''
271            i_name = '%s_%s' % (model._meta.db_table, self._digest(f.column))
272            output = [style.SQL_KEYWORD('CREATE INDEX') + ' ' +
273                style.SQL_TABLE(qn(truncate_name(i_name, self.connection.ops.max_name_length()))) + ' ' +
274                style.SQL_KEYWORD('ON') + ' ' +
275                style.SQL_TABLE(qn(model._meta.db_table)) + ' ' +
276                "(%s)" % style.SQL_FIELD(qn(f.column)) +
277                "%s;" % tablespace_sql]
278        else:
279            output = []
280        return output
281
282    def sql_destroy_model(self, model, references_to_delete, style):
283        "Return the DROP TABLE and restraint dropping statements for a single model"
284        if not model._meta.managed or model._meta.proxy:
285            return []
286        # Drop the table now
287        qn = self.connection.ops.quote_name
288        output = ['%s %s;' % (style.SQL_KEYWORD('DROP TABLE'),
289                              style.SQL_TABLE(qn(model._meta.db_table)))]
290        if model in references_to_delete:
291            output.extend(self.sql_remove_table_constraints(model, references_to_delete, style))
292
293        if model._meta.has_auto_field:
294            ds = self.connection.ops.drop_sequence_sql(model._meta.db_table)
295            if ds:
296                output.append(ds)
297        return output
298
299    def sql_remove_table_constraints(self, model, references_to_delete, style):
300        from django.db.backends.util import truncate_name
301
302        if not model._meta.managed or model._meta.proxy:
303            return []
304        output = []
305        qn = self.connection.ops.quote_name
306        for rel_class, f in references_to_delete[model]:
307            table = rel_class._meta.db_table
308            col = f.column
309            r_table = model._meta.db_table
310            r_col = model._meta.get_field(f.rel.field_name).column
311            r_name = '%s_refs_%s_%s' % (col, r_col, self._digest(table, r_table))
312            output.append('%s %s %s %s;' % \
313                (style.SQL_KEYWORD('ALTER TABLE'),
314                style.SQL_TABLE(qn(table)),
315                style.SQL_KEYWORD(self.connection.ops.drop_foreignkey_sql()),
316                style.SQL_FIELD(qn(truncate_name(r_name, self.connection.ops.max_name_length())))))
317        del references_to_delete[model]
318        return output
319
320    def sql_destroy_many_to_many(self, model, f, style):
321        "Returns the DROP TABLE statements for a single m2m field"
322        import warnings
323        warnings.warn(
324            'Database creation API for m2m tables has been deprecated. M2M models are now automatically generated',
325            PendingDeprecationWarning
326        )
327
328        qn = self.connection.ops.quote_name
329        output = []
330        if f.auto_created:
331            output.append("%s %s;" % (style.SQL_KEYWORD('DROP TABLE'),
332                style.SQL_TABLE(qn(f.m2m_db_table()))))
333            ds = self.connection.ops.drop_sequence_sql("%s_%s" % (model._meta.db_table, f.column))
334            if ds:
335                output.append(ds)
336        return output
337   
338    def _set_test_dict(self):
339        if "TEST_NAME" in self.connection.settings_dict:
340            self.connection.settings_dict["NAME"] = self.connection.settings_dict["TEST_NAME"]
341        if "TEST_USER" in self.connection.settings_dict:
342            self.connection.settings_dict['USER'] = self.connection.settings_dict["TEST_USER"]
343        if "TEST_PASSWORD" in self.connection.settings_dict:
344            self.connection.settings_dict['PASSWORD'] = self.connection.settings_dict["TEST_PASSWORD"]
345
346    def create_test_db(self, verbosity=1, autoclobber=False):
347        """
348        Creates a test database, prompting the user for confirmation if the
349        database already exists. Returns the name of the test database created.
350        """
351        if verbosity >= 1:
352            print "Creating test database '%s'..." % self.connection.alias
353
354        test_database_name = self._create_test_db(verbosity, autoclobber)
355
356        self.connection.close()
357        self.connection.settings_dict["NAME"] = test_database_name
358        can_rollback = self._rollback_works()
359        self.connection.settings_dict["SUPPORTS_TRANSACTIONS"] = can_rollback
360
361        call_command('syncdb', verbosity=verbosity, interactive=False, database=self.connection.alias)
362
363        if settings.CACHE_BACKEND.startswith('db://'):
364            from django.core.cache import parse_backend_uri
365            _, cache_name, _ = parse_backend_uri(settings.CACHE_BACKEND)
366            call_command('createcachetable', cache_name)
367
368        # Get a cursor (even though we don't need one yet). This has
369        # the side effect of initializing the test database.
370        cursor = self.connection.cursor()
371
372        return test_database_name
373
374    def _create_test_db(self, verbosity, autoclobber):
375        "Internal implementation - creates the test db tables."
376       
377        suffix = self.sql_table_creation_suffix()
378
379        if self.connection.settings_dict['TEST_NAME']:
380            test_database_name = self.connection.settings_dict['TEST_NAME']
381        else:
382            test_database_name = TEST_DATABASE_PREFIX + self.connection.settings_dict['NAME']
383
384        qn = self.connection.ops.quote_name
385
386        # Create the test database and connect to it. We need to autocommit
387        # if the database supports it because PostgreSQL doesn't allow
388        # CREATE/DROP DATABASE statements within transactions.
389        self._set_test_dict()
390        cursor = self.connection.cursor()
391        self.set_autocommit()
392
393        return test_database_name
394
395    def _rollback_works(self):
396        cursor = self.connection.cursor()
397        cursor.execute('CREATE TABLE ROLLBACK_TEST (X INT)')
398        self.connection._commit()
399        cursor.execute('INSERT INTO ROLLBACK_TEST (X) VALUES (8)')
400        self.connection._rollback()
401        cursor.execute('SELECT COUNT(X) FROM ROLLBACK_TEST')
402        count, = cursor.fetchone()
403        cursor.execute('DROP TABLE ROLLBACK_TEST')
404        self.connection._commit()
405        return count == 0
406
407    def destroy_test_db(self, old_database_name, verbosity=1):
408        """
409        Destroy a test database, prompting the user for confirmation if the
410        database already exists. Returns the name of the test database created.
411        """
412        if verbosity >= 1:
413            print "Destroying test database '%s'..." % self.connection.alias
414        self.connection.close()
415        test_database_name = self.connection.settings_dict['NAME']
416        self.connection.settings_dict['NAME'] = old_database_name
417        self._destroy_test_db(test_database_name, verbosity)
418
419    def _destroy_test_db(self, test_database_name, verbosity):
420        "Internal implementation - remove the test db tables."
421       
422        # Remove the test database to clean up after
423        # ourselves. Connect to the previous database (not the test database)
424        # to do so, because it's not allowed to delete a database while being
425        # connected to it.
426        self._set_test_dict()
427        cursor = self.connection.cursor()
428        self.set_autocommit()
429        time.sleep(1) # To avoid "database is being accessed by other users" errors.
430       
431        cursor.execute("""SELECT table_name FROM information_schema.tables WHERE table_schema='public'""")
432        rows = cursor.fetchall()
433        dropped_tables = []
434        not_dropped = []
435        for row in rows:
436            try:
437                cursor.execute('drop table `%s` cascade' % row[0])
438                dropped_tables.append(row[0])
439            except:
440                not_dropped.append(row[0])
441               
442        print 'Dropped tables: %s' % ', '.join(dropped_tables)
443        if not_dropped: print 'Error: Could not drop: %s' % ', '.join(not_dropped)
444       
445        #cursor.execute("DROP DATABASE %s" % self.connection.ops.quote_name(test_database_name))
446        self.connection.close()
447
448    def set_autocommit(self):
449        "Make sure a connection is in autocommit mode."
450        if hasattr(self.connection.connection, "autocommit"):
451            if callable(self.connection.connection.autocommit):
452                self.connection.connection.autocommit(True)
453            else:
454                self.connection.connection.autocommit = True
455        elif hasattr(self.connection.connection, "set_isolation_level"):
456            self.connection.connection.set_isolation_level(0)
457
458    def sql_table_creation_suffix(self):
459        "SQL to append to the end of the test table creation statements"
460        return ''
461
Back to Top