sqlmapper.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
#  sqlmapper.py
#  PythonSQLMapper
#
#  Copyright 2013-2026 Kenji Nishishiro. All rights reserved.
#  Written by Kenji Nishishiro <marvel@programmershigh.org>.
#

import re


class MappingError(Exception):
    pass


class DriverWarning(MappingError):
    pass


class DriverError(MappingError):
    pass


class DriverInterfaceError(DriverError):
    pass


class DriverDatabaseError(DriverError):
    pass


class DriverDataError(DriverDatabaseError):
    pass


class DriverOperationalError(DriverDatabaseError):
    pass


class DriverIntegrityError(DriverDatabaseError):
    pass


class DriverInternalError(DriverDatabaseError):
    pass


class DriverProgrammingError(DriverDatabaseError):
    pass


class DriverNotSupportedError(DriverDatabaseError):
    pass


class Result(object):
    pass


class Mapper(object):
    def __init__(self, driver, **params):
        self.driver = driver
        self.connection = None

        if self.driver.__name__ == "sqlite3":
            self.__cursor_params = {}
            self.__buffered_cursor_params = self.__cursor_params
            self.__place_holder = "?"
        elif self.driver.__name__ == "mysql.connector":
            self.__cursor_params = {"dictionary": True}
            self.__buffered_cursor_params = {"dictionary": True, "buffered": True}
            self.__place_holder = "%s"
        elif self.driver.__name__ == "MySQLdb":
            import MySQLdb.cursors

            self.__cursor_params = {"cursorclass": MySQLdb.cursors.SSDictCursor}
            self.__buffered_cursor_params = {"cursorclass": MySQLdb.cursors.DictCursor}
            self.__place_holder = "%s"
        elif self.driver.__name__ == "pymysql":
            import pymysql.cursors

            self.__cursor_params = {"cursor": pymysql.cursors.SSDictCursor}
            self.__buffered_cursor_params = {"cursor": pymysql.cursors.DictCursor}
            self.__place_holder = "%s"
        elif self.driver.__name__ == "psycopg2":
            import psycopg2.extras

            self.__cursor_params = {"cursor_factory": psycopg2.extras.RealDictCursor}
            self.__buffered_cursor_params = self.__cursor_params
            self.__place_holder = "%s"
        else:
            raise MappingError(
                f"Unsupported driver '{self.driver.__name__}'. Supported drivers: sqlite3, mysql.connector, "
                "MySQLdb, pymysql, psycopg2."
            )

        try:
            self.connection = self.driver.connect(**params)
        except Exception as error:
            mapped = self.__map_driver_error(error)
            if mapped is not None:
                raise mapped from error
            else:
                raise

        if self.driver.__name__ == "sqlite3":
            self.connection.row_factory = self.__sqlite3_dict_row_factory

    def close(self):
        try:
            if self.connection is not None:
                self.connection.close()
                self.connection = None
        except Exception as error:
            mapped = self.__map_driver_error(error)
            if mapped is not None:
                raise mapped from error
            else:
                raise

    def __del__(self):
        self.close()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.close()

    def select_one(self, sql, parameter=None, result_type=None):
        try:
            cursor = self.connection.cursor(**self.__cursor_params)
            try:
                cursor.execute(*self.__map_parameter(sql, parameter))
                rows = cursor.fetchmany(2)
                if len(rows) == 0:
                    return None
                elif len(rows) == 1:
                    return self.__create_result(row=rows[0], result_type=result_type)
                else:
                    raise MappingError("Expected exactly one row, but multiple rows were returned.")
            finally:
                cursor.close()
        except Exception as error:
            mapped = self.__map_driver_error(error)
            if mapped is not None:
                raise mapped from error
            else:
                raise

    returning_one = select_one

    def select_all(self, sql, parameter=None, result_type=None, array_size=1, buffered=True):
        try:
            if buffered:
                cursor = self.connection.cursor(**self.__buffered_cursor_params)
            else:
                cursor = self.connection.cursor(**self.__cursor_params)
            try:
                cursor.execute(*self.__map_parameter(sql, parameter))
                rows = cursor.fetchmany(array_size)
                while rows:
                    for row in rows:
                        yield self.__create_result(row=row, result_type=result_type)
                    rows = cursor.fetchmany(array_size)
            finally:
                cursor.close()
        except Exception as error:
            mapped = self.__map_driver_error(error)
            if mapped is not None:
                raise mapped from error
            else:
                raise

    returning_all = select_all

    def insert(self, sql, parameter=None):
        try:
            cursor = self.connection.cursor(**self.__cursor_params)
            try:
                cursor.execute(*self.__map_parameter(sql, parameter))
                return cursor.lastrowid
            finally:
                cursor.close()
        except Exception as error:
            mapped = self.__map_driver_error(error)
            if mapped is not None:
                raise mapped from error
            else:
                raise

    def update(self, sql, parameter=None):
        try:
            cursor = self.connection.cursor(**self.__cursor_params)
            try:
                cursor.execute(*self.__map_parameter(sql, parameter))
                return cursor.rowcount
            finally:
                cursor.close()
        except Exception as error:
            mapped = self.__map_driver_error(error)
            if mapped is not None:
                raise mapped from error
            else:
                raise

    delete = update

    def upsert(self, sql, parameter=None):
        try:
            cursor = self.connection.cursor(**self.__cursor_params)
            try:
                cursor.execute(*self.__map_parameter(sql, parameter))
                return cursor.rowcount, cursor.lastrowid
            finally:
                cursor.close()
        except Exception as error:
            mapped = self.__map_driver_error(error)
            if mapped is not None:
                raise mapped from error
            else:
                raise

    ignore = upsert

    def execute(self, sql, parameter=None):
        try:
            cursor = self.connection.cursor(**self.__cursor_params)
            try:
                cursor.execute(*self.__map_parameter(sql, parameter))
            finally:
                cursor.close()
        except Exception as error:
            mapped = self.__map_driver_error(error)
            if mapped is not None:
                raise mapped from error
            else:
                raise

    def commit(self):
        try:
            self.connection.commit()
        except Exception as error:
            mapped = self.__map_driver_error(error)
            if mapped is not None:
                raise mapped from error
            else:
                raise

    def rollback(self):
        try:
            self.connection.rollback()
        except Exception as error:
            mapped = self.__map_driver_error(error)
            if mapped is not None:
                raise mapped from error
            else:
                raise

    def __map_parameter(self, sql, parameter):
        represented_sql = ""
        parameters = ()
        start = 0
        for match in re.finditer("(?<!:):[a-zA-Z_][a-zA-Z0-9_]*", sql):
            represented_sql += sql[start : match.start()] + self.__place_holder
            start = match.end()
            parameters += (self.__get_variable(parameter, sql[match.start() + 1 : match.end()]),)
        represented_sql += sql[start:]
        return represented_sql, parameters

    def __map_driver_error(self, error):
        if isinstance(error, self.driver.NotSupportedError):
            return DriverNotSupportedError(*error.args)
        elif isinstance(error, self.driver.ProgrammingError):
            return DriverProgrammingError(*error.args)
        elif isinstance(error, self.driver.InternalError):
            return DriverInternalError(*error.args)
        elif isinstance(error, self.driver.IntegrityError):
            return DriverIntegrityError(*error.args)
        elif isinstance(error, self.driver.OperationalError):
            return DriverOperationalError(*error.args)
        elif isinstance(error, self.driver.DataError):
            return DriverDataError(*error.args)
        elif isinstance(error, self.driver.DatabaseError):
            return DriverDatabaseError(*error.args)
        elif isinstance(error, self.driver.InterfaceError):
            return DriverInterfaceError(*error.args)
        elif isinstance(error, self.driver.Error):
            return DriverError(*error.args)
        elif isinstance(error, self.driver.Warning):
            return DriverWarning(*error.args)
        else:
            return None

    @staticmethod
    def __sqlite3_dict_row_factory(cursor, row):
        fields = [column[0] for column in cursor.description]
        return {key: value for key, value in zip(fields, row)}

    @staticmethod
    def __get_variable(parameter, name):
        if isinstance(parameter, dict):
            try:
                return parameter[name]
            except KeyError:
                raise MappingError(
                    f"Bind variable '{name}' was not found in dict parameter. Available keys: {sorted(parameter.keys())}"
                )
        else:
            try:
                return getattr(parameter, name)
            except AttributeError:
                raise MappingError(
                    f"Bind variable '{name}' was not found in parameter object of type '{type(parameter).__name__}'."
                )

    @staticmethod
    def __create_result(row, result_type):
        if result_type is None:
            result = Result()
            for name in row:
                setattr(result, name, row[name])
            return result
        else:
            try:
                result = result_type()
            except TypeError:
                raise MappingError(f"Result type '{result_type}' must be instantiable without arguments.")
            for name in row:
                if hasattr(result, name):
                    setattr(result, name, row[name])
                else:
                    raise MappingError(f"Attribute '{name}' was not found in result_type '{result_type.__name__}'.")
            return result