# Path: hecatoncheir/QueryResult.py
# class QueryResult():
#     query = None
#     column_names = None
#     resultset = None
# 
#     def __init__(self, query):
#         self.query = query
#         self.column_names = []
#         self.resultset = []
# 
#     def __repr__(self):
#         return u'QueryResult:' + repr({'query': self.query,
#                                        'column': self.column_names,
#                                        'result': self.resultset})
# 
# Path: hecatoncheir/exception.py
# class DriverError(Exception):
#     source = None
# 
#     def __init__(self, value, source=None):
#         self.value = value
#         self.source = source
# 
#     def __str__(self):
#         return self.value
# 
# class InternalError(Exception):
#     source = None
#     query = None
# 
#     def __init__(self, value, query=None, source=None):
#         self.value = value
#         self.query = query
#         self.source = source
# 
#     def __str__(self):
#         return self.value
# 
# class QueryError(Exception):
#     source = None
#     query = None
# 
#     def __init__(self, value, query, source=None):
#         self.value = value
#         self.query = query
#         self.source = source
# 
#     def __str__(self):
#         return self.value
# 
# class QueryTimeout(Exception):
#     source = None
#     query = None
# 
#     def __init__(self, value, query, source=None):
#         self.value = value
#         self.query = query
#         self.source = source
# 
#     def __str__(self):
#         return self.value
# 
# Path: hecatoncheir/pgsql/PgDriver.py
# class PgDriver(DbDriverBase.DbDriverBase):
#     connstr = None
#     dbuser = None
#     dbpass = None
#     conn = None
#     driver = None
# 
#     def __init__(self, connstr, dbuser, dbpass):
#         self.connstr = connstr
#         self.dbuser = dbuser
#         self.dbpass = dbpass
# 
#         name = "psycopg2"
#         try:
#             self.driver = __import__(name, fromlist=[''])
#         except Exception as e:
#             raise DriverError(
#                 u"Could not load the driver module: %s" % name, source=e)
# 
#     def connect(self):
#         s = self.connstr + " user=" + self.dbuser + " password=" + self.dbpass
#         try:
#             self.conn = self.driver.connect(s)
#         except Exception as e:
#             raise DriverError(
#                 u"Could not connect to the server: %s" %
#                 e.args[0].split('\n')[0], source=e)
# 
#         return True
# 
#     def query_to_resultset(self, query, max_rows=10000, timeout=None):
#         """Build a QueryResult object from the query
# 
#         Args:
#             query (str): a query string to be executed.
#             max_rows (int): max rows which can be kept in a QueryResult object.
# 
#         Returns:
#             QueryResult: an object holding query, column names and result set.
#         """
#         assert query
#         assert isinstance(query, unicode)
#         log.trace('query_to_resultset: start query=%s' % query)
# 
#         res = QueryResult(query)
#         try:
#             if self.conn is None:
#                 self.connect()
# 
#             cur = self.conn.cursor()
#             if timeout and int(timeout) > 0:
#                 cur.execute('set statement_timeout to %d' % (timeout*1000))
# 
#             cur.execute(res.query)
# 
#             desc = []
#             for d in cur.description:
#                 desc.append(d[0])
#             res.column_names = deepcopy(tuple(desc))
# 
#             for i, r in enumerate(cur.fetchall()):
#                 # let's consider the memory size.
#                 if i > max_rows:
#                     raise InternalError(
#                         u'Exceeded the record limit (%d) for QueryResult.' %
#                         max_rows, query=query)
#                 res.resultset.append(deepcopy(r))
#             cur.close()
#         except InternalError as e:
#             raise e
#         except DriverError as e:
#             raise e
#         except Exception as e:
#             msg = unicode(e).split('\n')[0]
#             if msg == 'canceling statement due to statement timeout':
#                 raise QueryTimeout(
#                     "Query timeout: %s" % query,
#                     query=query, source=e)
#             raise QueryError(
#                 "Could not execute a query: %s" % msg,
#                 query=query, source=e)
#         finally:
#             if self.conn:
#                 self.conn.rollback()
# 
#         log.trace('query_to_resultset: end')
#         return res
# 
#     def disconnect(self):
#         if self.conn is None:
#             return False
# 
#         try:
#             self.conn.close()
#         except Exception as e:
#             raise DriverError(
#                 u"Could not disconnect from the server: %s" %
#                 e.args[0].split('\n')[0], source=e)
#         self.conn = None
#         return True


import os
import sys
import unittest
from hecatoncheir.QueryResult import QueryResult
from hecatoncheir.exception import DriverError, InternalError, QueryError, QueryTimeout
from hecatoncheir.pgsql import PgDriver


class TestPgDriver(unittest.TestCase):
    dbname = None
    dbuser = None
    dbpass = None

    def setUp(self):
        self.dbname = os.environ.get('PGDATABASE', 'dqwbtest')
        self.dbuser = os.environ.get('PGUSER', 'dqwbuser')
        self.dbpass = os.environ.get('PGPASSWORD', 'dqwbuser')

    def test_PgDriver_001(self):
        pg = PgDriver.PgDriver('a', 'b', 'c')
        self.assertTrue(pg is not None)
        self.assertEqual('a', pg.connstr)
        self.assertEqual('b', pg.dbuser)
        self.assertEqual('c', pg.dbpass)

    def test_connect_001(self):
        # connection success
        pg = PgDriver.PgDriver('host=/tmp dbname=%s' % self.dbname, self.dbuser, self.dbpass)
        try:
            pg.connect()
        except DriverError as e:
            self.fail()
        self.assertIsNotNone(pg.conn)

    def test_connect_002(self):
        # connection failure

# Path: hecatoncheir/QueryResult.py
# class QueryResult():
#     query = None
#     column_names = None
#     resultset = None
# 
#     def __init__(self, query):
#         self.query = query
#         self.column_names = []
#         self.resultset = []
# 
#     def __repr__(self):
#         return u'QueryResult:' + repr({'query': self.query,
#                                        'column': self.column_names,
#                                        'result': self.resultset})
# 
# Path: hecatoncheir/exception.py
# class DriverError(Exception):
#     source = None
# 
#     def __init__(self, value, source=None):
#         self.value = value
#         self.source = source
# 
#     def __str__(self):
#         return self.value
# 
# class InternalError(Exception):
#     source = None
#     query = None
# 
#     def __init__(self, value, query=None, source=None):
#         self.value = value
#         self.query = query
#         self.source = source
# 
#     def __str__(self):
#         return self.value
# 
# class QueryError(Exception):
#     source = None
#     query = None
# 
#     def __init__(self, value, query, source=None):
#         self.value = value
#         self.query = query
#         self.source = source
# 
#     def __str__(self):
#         return self.value
# 
# class QueryTimeout(Exception):
#     source = None
#     query = None
# 
#     def __init__(self, value, query, source=None):
#         self.value = value
#         self.query = query
#         self.source = source
# 
#     def __str__(self):
#         return self.value
# 
# Path: hecatoncheir/pgsql/PgDriver.py
# class PgDriver(DbDriverBase.DbDriverBase):
#     connstr = None
#     dbuser = None
#     dbpass = None
#     conn = None
#     driver = None
# 
#     def __init__(self, connstr, dbuser, dbpass):
#         self.connstr = connstr
#         self.dbuser = dbuser
#         self.dbpass = dbpass
# 
#         name = "psycopg2"
#         try:
#             self.driver = __import__(name, fromlist=[''])
#         except Exception as e:
#             raise DriverError(
#                 u"Could not load the driver module: %s" % name, source=e)
# 
#     def connect(self):
#         s = self.connstr + " user=" + self.dbuser + " password=" + self.dbpass
#         try:
#             self.conn = self.driver.connect(s)
#         except Exception as e:
#             raise DriverError(
#                 u"Could not connect to the server: %s" %
#                 e.args[0].split('\n')[0], source=e)
# 
#         return True
# 
#     def query_to_resultset(self, query, max_rows=10000, timeout=None):
#         """Build a QueryResult object from the query
# 
#         Args:
#             query (str): a query string to be executed.
#             max_rows (int): max rows which can be kept in a QueryResult object.
# 
#         Returns:
#             QueryResult: an object holding query, column names and result set.
#         """
#         assert query
#         assert isinstance(query, unicode)
#         log.trace('query_to_resultset: start query=%s' % query)
# 
#         res = QueryResult(query)
#         try:
#             if self.conn is None:
#                 self.connect()
# 
#             cur = self.conn.cursor()
#             if timeout and int(timeout) > 0:
#                 cur.execute('set statement_timeout to %d' % (timeout*1000))
# 
#             cur.execute(res.query)
# 
#             desc = []
#             for d in cur.description:
#                 desc.append(d[0])
#             res.column_names = deepcopy(tuple(desc))
# 
#             for i, r in enumerate(cur.fetchall()):
#                 # let's consider the memory size.
#                 if i > max_rows:
#                     raise InternalError(
#                         u'Exceeded the record limit (%d) for QueryResult.' %
#                         max_rows, query=query)
#                 res.resultset.append(deepcopy(r))
#             cur.close()
#         except InternalError as e:
#             raise e
#         except DriverError as e:
#             raise e
#         except Exception as e:
#             msg = unicode(e).split('\n')[0]
#             if msg == 'canceling statement due to statement timeout':
#                 raise QueryTimeout(
#                     "Query timeout: %s" % query,
#                     query=query, source=e)
#             raise QueryError(
#                 "Could not execute a query: %s" % msg,
#                 query=query, source=e)
#         finally:
#             if self.conn:
#                 self.conn.rollback()
# 
#         log.trace('query_to_resultset: end')
#         return res
# 
#     def disconnect(self):
#         if self.conn is None:
#             return False
# 
#         try:
#             self.conn.close()
#         except Exception as e:
#             raise DriverError(
#                 u"Could not disconnect from the server: %s" %
#                 e.args[0].split('\n')[0], source=e)
#         self.conn = None
#         return True


# Path: tests/testPgDriver.py
import os
import sys
import unittest
from hecatoncheir.QueryResult import QueryResult
from hecatoncheir.exception import DriverError, InternalError, QueryError, QueryTimeout
from hecatoncheir.pgsql import PgDriver



class TestPgDriver(unittest.TestCase):
    dbname = None
    dbuser = None
    dbpass = None

    def setUp(self):
        self.dbname = os.environ.get('PGDATABASE', 'dqwbtest')
        self.dbuser = os.environ.get('PGUSER', 'dqwbuser')
        self.dbpass = os.environ.get('PGPASSWORD', 'dqwbuser')

    def test_PgDriver_001(self):
        pg = PgDriver.PgDriver('a', 'b', 'c')
        self.assertTrue(pg is not None)
        self.assertEqual('a', pg.connstr)
        self.assertEqual('b', pg.dbuser)
        self.assertEqual('c', pg.dbpass)

    def test_connect_001(self):
        # connection success
        pg = PgDriver.PgDriver('host=/tmp dbname=%s' % self.dbname, self.dbuser, self.dbpass)
        try:
            pg.connect()
        except DriverError as e:
            self.fail()
        self.assertIsNotNone(pg.conn)

    def test_connect_002(self):
        # connection failure

        pg = PgDriver.PgDriver('host=/tmp dbname=%s' % self.dbname, "nosuchuser", '')