Python django.db.connection.connection() Examples

def ensure_mysql_connection_usable():
    """Ensure that MySQL connection is usable

    From: http://stackoverflow.com/questions/7835272/django-operationalerror-2006-mysql-server-has-gone-away
    """
    from django.db import connection, connections
    # MySQL is lazily connected to in Django.
    # connection.connection is None means
    # you have not connected to MySQL before
    if connection.connection and not connection.is_usable():
        # destroy the default MySQL connection
        # after this line, when you use ORM methods
        # Django will reconnect to the default MySQL
        #
        # Delete one database connection:
        # del connections._connections.default
        #
        # Delete all database connections
        databases = connections._connections.__dict__.keys()
        for database in databases:
            del connections._connections.__dict__[database] 

def _ensureConnection(self):
        # If connection is already made close it.
        from django.db import connection

        if connection.connection is not None:
            connection.close()

        # Loop forever until a connection can be made.
        while True:
            try:
                connection.ensure_connection()
            except Exception:
                log.err(
                    _why=(
                        "Error starting: "
                        "Connection to database cannot be established."
                    )
                )
                time.sleep(1)
            else:
                # Connection made, now close it.
                connection.close()
                break 

def connected():
    """Context manager that ensures we're connected to the database.

    If there is not yet a connection to the database, this will connect on
    entry and disconnect on exit. Preexisting connections will be left alone.

    If the preexisting connection is not usable it is closed and a new
    connection is made.
    """
    if connection.connection is None:
        connection.close_if_unusable_or_obsolete()
        connection.ensure_connection()
        try:
            yield
        finally:
            connection.close()
    elif connection.is_usable():
        yield
    else:
        # Connection is not usable, so we disconnect and reconnect. Since
        # the connection was previously connected we do not disconnect this
        # new connection.
        connection.close_if_unusable_or_obsolete()
        connection.ensure_connection()
        yield 

def with_connection(func):
    """Ensure that we're connected to the database before calling `func`.

    If there is not yet a connection to the database, this will connect before
    calling the decorated function, and then it will disconnect when done.
    Preexisting connections will be left alone.

    This can be important when using non-transactional advisory locks.
    """

    @wraps(func)
    def call_with_connection(*args, **kwargs):
        with connected():
            return func(*args, **kwargs)

    # For convenience, when introspecting for example, expose the original
    # function on the function we're returning.
    call_with_connection.func = func

    return call_with_connection 

def savepoint():
    """Context manager to wrap the code within a savepoint.

    This also enters a savepoint context for post-commit hooks, and so should
    always be used in preference to `transaction.atomic()` when only a
    savepoint is needed.

    If either a transaction or a savepoint within a transaction is what you
    want, use the `transactional` decorator.

    If you want a _decorator_ specifically, use the `transactional` decorator.

    If you want a _savepoint decorator_ specifically, write one, or adapt
    this to do it.

    """
    if connection.in_atomic_block:
        with post_commit_hooks.savepoint():
            with transaction.atomic():
                yield
    else:
        raise TransactionManagementError(
            "Savepoints cannot be created outside of a transaction."
        ) 

def is_locked(self):
        stmt = (
            "SELECT 1 FROM pg_locks, pg_database"
            " WHERE pg_locks.locktype = 'advisory'"
            "   AND pg_locks.classid = %s"
            "   AND pg_locks.objid = %s"
            # objsubid is 2 when using the 2-argument version of the
            # pg_advisory_* locking functions.
            "   AND pg_locks.objsubid = 2"
            "   AND pg_locks.granted"
            # Advisory locks are local to each database so we join to
            # pg_databases to discover the OID of the currrent database.
            "   AND pg_locks.database = pg_database.oid"
            "   AND pg_database.datname = current_database()"
        )
        with closing(connection.cursor()) as cursor:
            cursor.execute(stmt, self)
            return len(cursor.fetchall()) >= 1 

def listen(self):
        assert self.notify_channel, "You must set a notify channel in order to listen."
        with connection.cursor() as cur:
            cur.execute('LISTEN "{}";'.format(self.notify_channel)) 

def disable_all_database_connections():
    """Replace all connections in this thread with unusable stubs.

    Specifically, instances of :py:class:`~DisabledDatabaseConnection`.
    This should help prevent accidental use of the database from the
    reactor thread.

    Why?

    Database access means blocking IO, at least with the connections
    that Django hands out. While blocking IO isn't forbidden in the
    reactor thread, it ought to be avoided, because the reactor can't do
    anything else while it's happening, like handling other IO, or
    running delayed calls.

    Django's transaction and connection management code also assumes
    threads: it associates connections and transactions with the current
    thread, using threading.local. Using the database from the reactor
    thread is a recipe for intermingled transactions.
    """
    for alias in connections:
        connection = connections[alias]
        if type(connection) is not DisabledDatabaseConnection:
            connections[alias] = DisabledDatabaseConnection()
            connection.close() 

def test_with_client(self):
        with CaptureQueriesContext(connection) as captured_queries:
            self.client.get("/test_utils/get_person/%s/" % self.person_pk)
        self.assertEqual(len(captured_queries), 1)
        self.assertIn(self.person_pk, captured_queries[0]['sql'])

        with CaptureQueriesContext(connection) as captured_queries:
            self.client.get("/test_utils/get_person/%s/" % self.person_pk)
        self.assertEqual(len(captured_queries), 1)
        self.assertIn(self.person_pk, captured_queries[0]['sql'])

        with CaptureQueriesContext(connection) as captured_queries:
            self.client.get("/test_utils/get_person/%s/" % self.person_pk)
            self.client.get("/test_utils/get_person/%s/" % self.person_pk)
        self.assertEqual(len(captured_queries), 2)
        self.assertIn(self.person_pk, captured_queries[0]['sql'])
        self.assertIn(self.person_pk, captured_queries[1]['sql']) 

def test_ignores_connection_configuration_queries(self):
        real_ensure_connection = connection.ensure_connection
        connection.close()

        def make_configuration_query():
            is_opening_connection = connection.connection is None
            real_ensure_connection()

            if is_opening_connection:
                # Avoid infinite recursion. Creating a cursor calls
                # ensure_connection() which is currently mocked by this method.
                connection.cursor().execute('SELECT 1' + connection.features.bare_select_suffix)

        ensure_connection = 'django.db.backends.base.base.BaseDatabaseWrapper.ensure_connection'
        with mock.patch(ensure_connection, side_effect=make_configuration_query):
            with self.assertNumQueries(1):
                list(Car.objects.all()) 

def test_disconnects_and_reconnects_if_not_usable(self):
        connection.ensure_connection()
        preexisting_connection = connection.connection

        connection.errors_occurred = True
        self.patch(connection, "is_usable").return_value = False

        self.assertThat(connection.connection, Not(Is(None)))
        with orm.connected():
            self.assertThat(
                connection.connection, Not(Is(preexisting_connection))
            )
            self.assertThat(connection.connection, Not(Is(None)))

        self.assertThat(connection.connection, Not(Is(preexisting_connection)))
        self.assertThat(connection.connection, Not(Is(None))) 

def test_leaves_preexisting_connections_open(self):
        # Ensure there's a database connection to begin with.
        connection.ensure_connection()

        # No transaction has been entered (what Django calls an atomic block),
        # but the connection has been established.
        self.assertFalse(connection.in_atomic_block)
        self.expectThat(connection.connection, Not(Is(None)))

        # Call a function via the `transactional` decorator.
        decorated_function = orm.transactional(lambda: None)
        decorated_function()

        # After the decorated function has returned the transaction has
        # been exited, but the preexisting connection remains open.
        self.assertFalse(connection.in_atomic_block)
        self.expectThat(connection.connection, Not(Is(None))) 

def test_closes_connections_only_when_leaving_atomic_block(self):
        # Close the database connection to begin with.
        connection.close()
        self.expectThat(connection.connection, Is(None))

        @orm.transactional
        def inner():
            # We're inside a `transactional` context here.
            self.expectThat(connection.connection, Not(Is(None)))
            return "inner"

        @orm.transactional
        def outer():
            # We're inside a `transactional` context here too.
            self.expectThat(connection.connection, Not(Is(None)))
            # Call `inner`, thus nesting `transactional` contexts.
            return "outer > " + inner()

        self.assertEqual("outer > inner", outer())
        # The connection has been closed.
        self.expectThat(connection.connection, Is(None)) 

def test_ignores_connection_configuration_queries(self):
        real_ensure_connection = connection.ensure_connection
        connection.close()

        def make_configuration_query():
            is_opening_connection = connection.connection is None
            real_ensure_connection()

            if is_opening_connection:
                # Avoid infinite recursion. Creating a cursor calls
                # ensure_connection() which is currently mocked by this method.
                connection.cursor().execute('SELECT 1' + connection.features.bare_select_suffix)

        ensure_connection = 'django.db.backends.base.base.BaseDatabaseWrapper.ensure_connection'
        with mock.patch(ensure_connection, side_effect=make_configuration_query):
            with self.assertNumQueries(1):
                list(Car.objects.all()) 

def test_logs_all_queries_made_by_func(self):
        def query_func():
            return list(Node.objects.all())

        mock_print = Mock()
        wrapped = count_queries(mock_print)(query_func)
        wrapped()

        query_time = sum(
            [float(query.get("time", 0)) for query in connection.queries]
        )
        self.assertThat(
            mock_print,
            MockCalledOnceWith(
                "[QUERIES] query_func executed 1 queries in %s seconds"
                % query_time
            ),
        ) 

def get_select_query_txt(
    table_name: str,
    column_names: Optional[List[str]] = None,
    filter_formula: Optional[Dict] = None,
    filter_pairs: Optional[Mapping] = None,
) -> Tuple[str, List[Any]]:
    """Calculate the text representation of a query to select table subset.

    :param table_name: Table to query
    :param column_names: list of columns to consider or None to consider all
    :param filter_formula: Text filter expression
    :param filter_pairs: Dictionary of key/value pairs.
    :return: (sql query, sql params)
    """
    # invoke get_select_query and transform into string
    query_str, fields = get_select_query(
        table_name,
        column_names=column_names,
        filter_formula=filter_formula,
        filter_pairs=filter_pairs,
    )

    return query_str.as_string(connection.connection), fields 

def clear_firmware_data(self):
        print("Clearing firmware_flash.Firmware objects...")
        constraint_check = connection.disable_constraint_checking()
        firmware_flash.models.Firmware.objects.all().delete()

        print("Clearing firmware_flash.Board objects...")
        constraint_check = connection.disable_constraint_checking()
        firmware_flash.models.Board.objects.all().delete()

        print("Clearing firmware_flash.DeviceFamily objects...")
        constraint_check = connection.disable_constraint_checking()
        firmware_flash.models.DeviceFamily.objects.all().delete()

        print("Clearing firmware_flash.Project objects...")
        constraint_check = connection.disable_constraint_checking()
        firmware_flash.models.Project.objects.all().delete()


        print("Done clearing firmware_flash objects...") 

def close_connection():
    """Closes the connection if we are not in an atomic block.

    The connection should never be closed if we are in an atomic block, as
    happens when running tests as part of a django TestCase. Otherwise, closing
    the connection is important to avoid a connection time out after long actions.
    Django does not automatically refresh a connection which has been closed
    due to idleness (this normally happens in the request start/finish part
    of a webapp's lifecycle, which this process does not have), so we must
    do it ourselves if the connection goes idle due to stuff taking a really
    long time.

    source: http://stackoverflow.com/a/39322632/865091
    """
    from django.db import connection
    if not connection.in_atomic_block:
        connection.close() 

def test_ensures_connection(self):
        with orm.connected():
            self.assertThat(connection.connection, Not(Is(None))) 

def checkDatabaseUse(self):
        """Enforce `database_use_permitted`."""
        if self.database_use_possible and not self.database_use_permitted:
            from django.db import connection

            self.expectThat(
                connection.connection,
                testtools.matchers.Is(None),
                "Test policy forbids use of the database.",
            )
            connection.close() 

def test_resets_queries_between_calls(self):
        def query_func():
            return list(Node.objects.all())

        mock_print = Mock()
        wrapped = count_queries(mock_print)(query_func)

        # First call.
        wrapped()
        query_time_one = sum(
            [float(query.get("time", 0)) for query in connection.queries]
        )

        # Second call.
        wrapped()
        query_time_two = sum(
            [float(query.get("time", 0)) for query in connection.queries]
        )

        # Print called twice.
        self.assertThat(
            mock_print,
            MockCallsMatch(
                call(
                    "[QUERIES] query_func executed 1 queries in %s seconds"
                    % query_time_one
                ),
                call(
                    "[QUERIES] query_func executed 1 queries in %s seconds"
                    % query_time_two
                ),
            ),
        ) 

def assertClosed(self, alias):
        self.assertThat(connections[alias].connection, Is(None)) 

def in_transaction(_connection=None):
    """Is `_connection` in the midst of a transaction?

    This only enquires as to Django's perspective on the situation. It does
    not actually check that the database agrees with Django.

    :return: bool
    """
    if _connection is None:
        return connection.in_atomic_block
    else:
        return _connection.in_atomic_block 

def assertOpen(self, alias):
        self.assertThat(connections[alias].connection, Not(Is(None))) 

def __enter__(self):
        """Assert that no connections are yet open."""
        for alias in connections:
            if connections[alias].connection is not None:
                raise AssertionError("Connection %s is open." % (alias,)) 

def test_exit_closes_open_connections(self):
        self.addCleanup(connection.close)
        connection.ensure_connection()
        self.assertThat(connection.connection, Not(Is(None)))
        context = ExclusivelyConnected()
        context.__exit__()
        self.assertThat(connection.connection, Is(None)) 

def count_queries(log_func):
    """Decorator that will count the queries and call `log_func` with the log
    message.
    """

    def wrapper(func):
        def inner_wrapper(*args, **kwargs):
            # Reset the queries count before performing the function.
            reset_queries()

            # Perform the work that will create queries.
            result = func(*args, **kwargs)

            # Calculate the query_time and log the number and time.
            query_time = sum(
                [float(query.get("time", 0)) for query in connection.queries]
            )
            log_func(
                "[QUERIES] %s executed %s queries in %s seconds"
                % (func.__name__, len(connection.queries), query_time)
            )

            # Log all the queries if requested.
            if getattr(func, "__log_sql_calls__", False) or getattr(
                settings, "DEBUG_QUERIES_LOG_ALL", False
            ):
                log_func("[QUERIES] === Start SQL Log: %s ===" % func.__name__)
                for query in connection.queries:
                    log_func("[QUERIES] %s" % query.get("sql"))
                log_func("[QUERIES] === End SQL Log: %s ===" % func.__name__)

            return result

        return inner_wrapper

    return wrapper 

def test_opens_and_closes_connection_when_no_preexisting_connection(self):
        connection.close()

        self.assertThat(connection.connection, Is(None))
        with orm.connected():
            self.assertThat(connection.connection, Not(Is(None)))
        self.assertThat(connection.connection, Is(None)) 

def test_enter_blows_up_if_there_are_open_connections(self):
        self.addCleanup(connection.close)
        connection.ensure_connection()
        context = ExclusivelyConnected()
        self.assertRaises(AssertionError, context.__enter__) 

def test_exit_removes_block_on_database_connections(self):
        with TotallyDisconnected():
            self.assertRaises(RuntimeError, getattr, connection, "connect")
        connection.ensure_connection()