commons_database.py

#!/usr/bin/python
# -*- coding: utf-8 -*-

import os
import time

import pymysql


class DB:
    """
    Classe para fazer consultas ao banco de dados
    """

    def connect(self):
        username = os.environ.get('DB_USERNAME', None)
        password = os.environ.get('DB_PASSWORD', None)
        host = os.environ.get('DB_HOST', 'commonswiki.analytics.db.svc.eqiad.wmflabs')
        self.conn = pymysql.connect(
            db='commonswiki_p',
            host=host,
            user=username,
            passwd=password,
            read_default_file=os.path.expanduser('~/replica.my.cnf'),
            read_timeout=30, charset='utf8', use_unicode=True)
        self.conn.ping(True)

    def _query(self, *sql):
        with self.conn.cursor() as cursor:
            cursor.execute(*sql)
            return cursor.fetchall()

    def query(self, *sql):
        """
        Tenta fazer a consulta, reconecta até 10 vezes até conseguir
        """
        loops = 0
        self.connect()
        while True:
            try:
                return self._query(*sql)
            except (AttributeError, pymysql.err.OperationalError):
                if loops < 10:
                    loops += 1
                    print 'Erro no DB, esperando %ds antes de tentar de novo' % loops
                    time.sleep(loops)
                else:
                    return self._query(*sql)
                    break
            else:
                print "Uncaught exception when running query"
                print sql
                break
        self.close_connection()

    def close_connection(self):
        self.conn.close()