from functools import lru_cache

import pymonetdb


class MonetMemoize:

    def __init__(self, funcs=None, maxsize=128):
        self.maxsize = maxsize
        self.type_map = {
            'py_date': pymonetdb.sql.types.DATE,
            'py_time': pymonetdb.sql.types.TIME,
            'py_timestamp': pymonetdb.sql.types.TIMESTAMP,
            'py_timetz': pymonetdb.sql.types.TIMETZ,
            'py_timestamptz': pymonetdb.sql.types.TIMESTAMPTZ,
        }
        if funcs is None:
            self.funcs = tuple(self.type_map.keys())
        else:
            self.funcs = funcs

    def __enter__(self):
        for func in self.funcs:
            old_func = getattr(pymonetdb.sql.pythonize, func)
            monet_type = self.type_map[func]
            lru_func = lru_cache(maxsize=self.maxsize)(old_func)
            pymonetdb.sql.pythonize.mapping[monet_type] = lru_func

    def __exit__(self, *args, **kwargs):
        for func in self.funcs:
            old_func = getattr(pymonetdb.sql.pythonize, func)
            monet_type = self.type_map[func]
            pymonetdb.sql.pythonize.mapping[monet_type] = old_func


def main():
    connection = pymonetdb.connect(username='username', password='password',
                                   hostname='hostname', database='database')
    cursor = connection.cursor()

    with MonetMemoize():
        cursor.execute('select * from test;')
        print(len(cursor.fetchall()))


if __name__ == '__main__':
    main()