From aa76fce9aaa22ab11a7b472afc9ba0f7b4235065 Mon Sep 17 00:00:00 2001
From: Peng Yong <ppyy@pubyun.com>
Date: Sun, 11 Mar 2012 23:44:15 +0800
Subject: [PATCH] check connection in Listener. refer to Bug #943031

Change-Id: I9fd064d7a9e953ea7d958982c483ccc358dbf7fa
---
 Authors                          |  1 +
 melange/db/sqlalchemy/session.py | 33 ++++++++++++++++++++++++++++++--
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/Authors b/Authors
index e7a5bd8d..b3d29f6c 100644
--- a/Authors
+++ b/Authors
@@ -97,6 +97,7 @@ Naveed Massjouni <naveedm9@gmail.com>
 Nikolay Sokolov <nsokolov@griddynamics.com>
 Nirmal Ranganathan <nirmal.ranganathan@rackspace.com>
 Paul Voccio <paul@openstack.org>
+Peng Yong <ppyy@pubyun.com>
 Rajaram Mallya <rajarammallya@gmail.com>
 Renuka Apte <renuka.apte@citrix.com>
 Ricardo Carrillo Cruz <emaildericky@gmail.com>
diff --git a/melange/db/sqlalchemy/session.py b/melange/db/sqlalchemy/session.py
index 0156dc18..44ad987a 100644
--- a/melange/db/sqlalchemy/session.py
+++ b/melange/db/sqlalchemy/session.py
@@ -17,8 +17,10 @@
 
 import contextlib
 import logging
+import sqlalchemy as sql
 from sqlalchemy import create_engine
 from sqlalchemy import MetaData
+from sqlalchemy.exc import DisconnectionError
 from sqlalchemy.orm import sessionmaker
 
 from melange import ipam
@@ -53,17 +55,44 @@ def configure_sqlalchemy_log(options):
         logger.setLevel(logging.INFO)
 
 
+class MySQLPingListener(object):
+
+    """
+    Ensures that MySQL connections checked out of the
+    pool are alive.
+
+    Borrowed from:
+    http://groups.google.com/group/sqlalchemy/msg/a4ce563d802c929f
+    """
+
+    def checkout(self, dbapi_con, con_record, con_proxy):
+        try:
+            dbapi_con.cursor().execute('select 1')
+        except dbapi_con.OperationalError, ex:
+            if ex.args[0] in (2006, 2013, 2014, 2045, 2055):
+                LOG.warn('Got mysql server has gone away: %s', ex)
+                raise DisconnectionError("Database server went away")
+            else:
+                raise
+
+
 def _create_engine(options):
+    connection_dict = sql.engine.url.make_url(options['sql_connection'])
     engine_args = {
-        "pool_recycle": config.get_option(options,
+        'pool_recycle': config.get_option(options,
                                           'sql_idle_timeout',
                                           type='int',
                                           default=3600),
-        "echo": config.get_option(options,
+        'echo': config.get_option(options,
                                   'sql_query_log',
                                   type='bool',
                                   default=False),
+        'convert_unicode': True,
     }
+
+    if 'mysql' in connection_dict.drivername:
+        engine_args['listeners'] = [MySQLPingListener()]
+
     LOG.info("Creating SQLAlchemy engine with args: %s" % engine_args)
     return create_engine(options['sql_connection'], **engine_args)