Home | History | Annotate | Download | only in afe
      1 from django.db.backends.mysql.base import DatabaseCreation as MySQLCreation
      2 from django.db.backends.mysql.base import DatabaseOperations as MySQLOperations
      3 from django.db.backends.mysql.base import DatabaseWrapper as MySQLDatabaseWrapper
      4 from django.db.backends.mysql.base import DatabaseIntrospection as MySQLIntrospection
      5 
      6 try:
      7     import MySQLdb as Database
      8 except ImportError, e:
      9     from django.core.exceptions import ImproperlyConfigured
     10     raise ImproperlyConfigured("Error loading MySQLdb module: %s" % e)
     11 
     12 
     13 class DatabaseOperations(MySQLOperations):
     14     """Custom database backend wrapper."""
     15     compiler_module = "autotest_lib.frontend.db.backends.afe.compiler"
     16 
     17 
     18 class DatabaseWrapper(MySQLDatabaseWrapper):
     19     """Custom database backend wrapper."""
     20 
     21     def __init__(self, *args, **kwargs):
     22         self.connection = None
     23         super(DatabaseWrapper, self).__init__(*args, **kwargs)
     24         self.creation = MySQLCreation(self)
     25         try:
     26             self.ops = DatabaseOperations()
     27         except TypeError:
     28             self.ops = DatabaseOperations(connection=kwargs.get('connection'))
     29         self.introspection = MySQLIntrospection(self)
     30 
     31     def _valid_connection(self):
     32         if self.connection is not None:
     33             if self.connection.open:
     34                 try:
     35                     self.connection.ping()
     36                     return True
     37                 except Database.DatabaseError:
     38                     self.connection.close()
     39                     self.connection = None
     40         return False
     41