OwlCyberSecurity - MANAGER

Edit File: database.py

##
# Copyright (c) 2009-2017 Apple Inc. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
##

"""
Generic ADAPI database access object.
"""

__all__ = [
    "AbstractADBAPIDatabase",
]

import threading

try:
    from txdav.base.datastore.subpostgres import postgres
except ImportError:
    postgres = None

from twisted.enterprise.adbapi import ConnectionPool
from twisted.internet.defer import inlineCallbacks, returnValue
from twisted.python.threadpool import ThreadPool

from twext.python.log import Logger

from twistedcaldav.config import ConfigurationError

log = Logger()

class ConnectionCloseThread(threading.Thread):
    """
    An L{Thread} that closes its DB connection when it has finished running.
    """

def __init__(self, group=None, target=None, name=None,
                 args=(), kwargs=None, verbose=None):

if target is not None:
            self._realTarget = target
            target = self.targetWithConnectionClose
        super(ConnectionCloseThread, self).__init__(group=group, target=target, name=name, args=args, kwargs=kwargs, verbose=verbose)

def targetWithConnectionClose(self, *args, **kwargs):
        self._realTarget(*args, **kwargs)
        if hasattr(self, "_db_close"):
            self._db_close()

class ConnectionClosingThreadPool(ThreadPool):
    """
    A L{ThreadPool} that closes connections for each worker thread when stopped.
    """

threadFactory = ConnectionCloseThread

def stop(self):
        for tid, conn in self.pool.connections.items():
            for thread in self.threads:
                if thread.ident == tid:
                    thread._db_close = lambda: self.pool.disconnect(conn)
        ThreadPool.stop(self)

class AbstractADBAPIDatabase(object):
    """
    A generic SQL database.
    """

def __init__(self, dbID, dbapiName, dbapiArgs, persistent, **kwargs):
        """

@param persistent: C{True} if the data in the DB must be perserved during upgrades,
            C{False} if the DB data can be re-created from an external source.
        @type persistent: bool
        """
        self.dbID = dbID
        self.dbapiName = dbapiName
        self.dbapiArgs = dbapiArgs
        self.dbapikwargs = kwargs

self.persistent = persistent

self.initialized = False

def __repr__(self):
        return "<%s %r>" % (self.__class__.__name__, self.pool)

@inlineCallbacks
    def open(self):
        """
        Access the underlying database.
        @return: a db2 connection object for this index's underlying data store.
        """
        if not self.initialized:

self.pool = ConnectionPool(self.dbapiName, *self.dbapiArgs, **self.dbapikwargs)

# sqlite3 is not thread safe which means we have to close the sqlite3 connections in the same thread that
            # opened them. We need a special thread pool class that has a thread worker function that does a close
            # when a thread is closed.
            if self.dbapiName == "sqlite3":
                self.pool.threadpool.stop()
                self.pool.threadpool = ConnectionClosingThreadPool(1, 1)
                self.pool.threadpool.start()
                self.pool.threadpool.pool = self.pool

#
            # Set up the schema
            #
            # Create CALDAV table if needed

try:
                test = (yield self._test_schema_table())
                if test:
                    version = (yield self._db_value_for_sql("select VALUE from CALDAV where KEY = 'SCHEMA_VERSION'"))
                    dbtype = (yield self._db_value_for_sql("select VALUE from CALDAV where KEY = 'TYPE'"))

if (version != self._db_version()) or (dbtype != self._db_type()):

if dbtype != self._db_type():
                            log.error(
                                "Database {db} has different type ({t1} vs. {t2})",
                                db=self.dbID, t1=dbtype, t2=self._db_type()
                            )

# Delete this index and start over
                            yield self._db_remove()
                            yield self._db_init()

elif version != self._db_version():
                            log.error(
                                "Database {db} has different schema (v.{v1} vs. v.{v2})",
                                db=self.dbID, v1=version, v2=self._db_version()
                            )

# Upgrade the DB
                            yield self._db_upgrade(version)

else:
                    yield self._db_init()
                self.initialized = True
            except:
                # Clean up upon error so we don't end up leaking threads
                self.pool.close()
                self.pool = None
                raise

def close(self):

if self.initialized:
            try:
                self.pool.close()
            except Exception, e:
                log.error("Error whilst closing connection pool: {ex}", ex=e)
            self.pool = None
            self.initialized = False

@inlineCallbacks
    def clean(self):