velocity-python 0.0.35__py3-none-any.whl → 0.0.65__py3-none-any.whl

velocity/db/core/engine.py

@@ -1,11 +1,17 @@
+import inspect
+import sys
+import re
+import traceback
+from functools import wraps
 from velocity.db import exceptions
 from velocity.db.core.transaction import Transaction
 
-from functools import wraps
-import inspect, sys, re, traceback
 
+class Engine:
+    """
+    Encapsulates driver config, connection logic, error handling, and transaction decoration.
+    """
 
-class Engine(object):
     MAX_RETRIES = 100
 
     def __init__(self, driver, config, sql):
@@ -14,16 +20,12 @@ class Engine(object):
         self.__driver = driver
 
     def __str__(self):
-        return """[{}] engine({})""".format(self.sql.server, self.config)
+        return f"[{self.sql.server}] engine({self.config})"
 
     def connect(self):
         """
         Connects to the database and returns the connection object.
-
-        If the database is missing, it creates the database and then connects to it.
-
-        Returns:
-            conn: The connection object to the database.
+        If the database is missing, tries to create it, then reconnect.
         """
         try:
             conn = self.__connect()
@@ -36,158 +38,132 @@ class Engine(object):
 
     def __connect(self):
         """
-        Connects to the database using the provided configuration.
-
-        Returns:
-            A connection object representing the connection to the database.
-
-        Raises:
-            Exception: If the configuration parameter is not handled properly.
-            ProcessError is called to handle other exceptions.
+        Internal connection logic, raising suitable exceptions on error.
         """
         try:
             if isinstance(self.config, dict):
                 return self.driver.connect(**self.config)
-            elif isinstance(self.config, (tuple, list)):
+            if isinstance(self.config, (tuple, list)):
                 return self.driver.connect(*self.config)
-            elif isinstance(self.config, str):
+            if isinstance(self.config, str):
                 return self.driver.connect(self.config)
-            else:
-                raise Exception("Unhandled configuration parameter")
+            raise Exception("Unhandled configuration parameter.")
         except:
-            self.ProcessError()
+            self.process_error()
 
     def transaction(self, func_or_cls=None):
         """
-        Decorator for defining a transaction. Use this to wrap a function, method, or class to automatically
-            start a transaction if necessary. If the function, method or class is called with a `tx` keyword argument,
-            it will use that transaction object instead of creating a new one. If the function, method or class
-            is called with a `tx` positional argument, it will use that transaction object instead of creating a new one.
-            If the function, method or class is called with a `tx` positional argument and a `tx` keyword argument, it will use the positional
-            argument and ignore the keyword argument. If the function, method or class is called without a `tx` argument,
-            it will create a new transaction object and use that.
-        Args:
-            func_or_cls: The function or class to be decorated.
-
-        Returns:
-            If `func_or_cls` is a function or method, returns a wrapped version of the function or method that
-            automatically starts a transaction if necessary. If `func_or_cls` is a class, returns a subclass of
-            `func_or_cls` that wraps all its methods with the transaction decorator.
-
-
-
-        If `func_or_cls` is not provided, returns a new `Transaction` object associated with the engine.
-        """
-        # If you are having trouble passing TWO transaction objects, for
-        # example as a source database to draw data from, pass the second as
-        # a keyword. For example:
-        #       @engine.transaction
-        #       def function(tx, src=src) <-- pass second as a kwd and tx will populate correctly.
-        #           ...
-        #
-        engine = self
+        Decorator that provides a Transaction. If `tx` is passed in, uses it; otherwise, creates a new one.
+        May also be used to decorate a class, in which case all methods are wrapped in a transaction if they accept `tx`.
+        With no arguments, returns a new Transaction directly.
+        """
+        # print("Transaction", func_or_cls.__name__, type(func_or_cls))
+
+        if func_or_cls is None:
+            return Transaction(self)
+
+        if isinstance(func_or_cls, classmethod):
+            return classmethod(self.transaction(func_or_cls.__func__))
+
         if inspect.isfunction(func_or_cls) or inspect.ismethod(func_or_cls):
+            names = list(inspect.signature(func_or_cls).parameters.keys())
+            # print(func_or_cls.__name__, names)
+            if "_tx" in names:
+                raise NameError(
+                    f"In function {func_or_cls.__name__}, '_tx' is not allowed as a parameter."
+                )
 
             @wraps(func_or_cls)
-            def NewFunction(*args, **kwds):
+            def new_function(*args, **kwds):
                 tx = None
                 names = list(inspect.signature(func_or_cls).parameters.keys())
-                if "_tx" in names:
-                    raise NameError(
-                        f"In function named `{func_or_cls.__name__}` You may not name a paramater `_tx`"
-                    )
+
+                # print("inside", func_or_cls.__name__)
+                # print(names)
+                # print(args, kwds)
+
                 if "tx" not in names:
+                    # The function doesn't even declare a `tx` parameter, so run normally.
                     return func_or_cls(*args, **kwds)
-                elif "tx" in kwds:
+
+                if "tx" in kwds:
                     if isinstance(kwds["tx"], Transaction):
                         tx = kwds["tx"]
                     else:
                         raise TypeError(
-                            f"In function named `{func_or_cls.__name__}` keyword `tx` must be a Transaction object"
+                            f"In function {func_or_cls.__name__}, keyword argument `tx` must be a Transaction object."
                         )
-                elif "tx" in names:
+                else:
+                    # Might be in positional args
                     pos = names.index("tx")
-                    if len(args) > pos and isinstance(args[pos], Transaction):
-                        tx = args[pos]
+                    if len(args) > pos:
+                        if isinstance(args[pos], Transaction):
+                            tx = args[pos]
+
                 if tx:
-                    return engine.exec_function(func_or_cls, tx, *args, **kwds)
-                else:
-                    with Transaction(engine) as tx:
-                        pos = names.index("tx")
-                        args = list(args)
-                        args.insert(pos, tx)
-                        args = tuple(args)
-                        return engine.exec_function(func_or_cls, tx, *args, **kwds)
-
-            return NewFunction
-        elif inspect.isclass(func_or_cls):
-
-            class NewCls(func_or_cls):
-                def __getattribute__(self, key):
-                    attr = super(NewCls, self).__getattribute__(key)
-                    if key in ["start_response"]:
-                        return attr
-                    if inspect.ismethod(attr):
-                        return engine.transaction(attr)
-                    return attr
+                    return self.exec_function(func_or_cls, tx, *args, **kwds)
+
+                with Transaction(self) as local_tx:
+                    pos = names.index("tx")
+                    new_args = args[:pos] + (local_tx,) + args[pos:]
+                    return self.exec_function(func_or_cls, local_tx, *new_args, **kwds)
+
+            return new_function
+
+        if inspect.isclass(func_or_cls):
+
+            NewCls = type(func_or_cls.__name__, (func_or_cls,), {})
+
+            for attr_name in dir(func_or_cls):
+                # Optionally skip special methods
+                if attr_name.startswith("__") and attr_name.endswith("__"):
+                    continue
+
+                attr = getattr(func_or_cls, attr_name)
+
+                if callable(attr):
+                    setattr(NewCls, attr_name, self.transaction(attr))
 
             return NewCls
 
-        return Transaction(engine)
+        return Transaction(self)
 
     def exec_function(self, function, _tx, *args, **kwds):
         """
-        Executes the given function with the provided arguments and keyword arguments.
-        If there is no transaction object, it executes the function without a transaction.
-
-        If there is a transaction object, it executes the function within the transaction.
-        If the function raises a `DbRetryTransaction` exception, it rolls back the transaction and retries.
-        If the function raises a `DbLockTimeoutError` exception, it rolls back the transaction and retries.
-        If any other exception occurs, it raises the exception.
-
-        Args:
-            function: The function to be executed.
-            tx: The transaction object to be passed to the function.
-            *args: Positional arguments to be passed to the function.
-            **kwds: Keyword arguments to be passed to the function.
-
-        Returns:
-            The result of the function execution.
-
-        Raises:
-            DbRetryTransaction: If the maximum number of retries is exceeded.
-            DbLockTimeoutError: If the maximum number of retries is exceeded.
-            Any other exception raised by the function.
-        """
-        retry_count = 0
-        tmout_count = 0
-        if _tx is None:
-            return function(*args, **kwds)
-        else:
-            while True:
-                try:
-                    return function(*args, **kwds)
-                except exceptions.DbRetryTransaction as e:
-                    if e.args and e.args[0]:
-                        print(e)
-                        print("**Retry Transaction. Rollback and start over")
+        Executes the given function inside the transaction `_tx`.
+        Retries if it raises DbRetryTransaction or DbLockTimeoutError, up to MAX_RETRIES times.
+        """
+        depth = getattr(_tx, "_exec_function_depth", 0)
+        setattr(_tx, "_exec_function_depth", depth + 1)
+
+        try:
+            if depth > 0:
+                # Not top-level. Just call the function.
+                return function(*args, **kwds)
+            else:
+                retry_count = 0
+                lock_timeout_count = 0
+                while True:
+                    try:
+                        return function(*args, **kwds)
+                    except exceptions.DbRetryTransaction as e:
+                        retry_count += 1
+                        if retry_count > self.MAX_RETRIES:
+                            raise
+                        _tx.rollback()
+                    except exceptions.DbLockTimeoutError as e:
+                        lock_timeout_count += 1
+                        if lock_timeout_count > self.MAX_RETRIES:
+                            raise
                         _tx.rollback()
                         continue
-                    retry_count += 1
-                    if retry_count > self.MAX_RETRIES:
+                    except:
                         raise
-                    print("**Retry Transaction. Rollback and start over")
-                    _tx.rollback()
-                    continue
-                except exceptions.DbLockTimeoutError:
-                    tmout_count += 1
-                    if tmout_count > self.MAX_RETRIES:
-                        raise
-                    print("**DbLockTimeoutError. Rollback and start over")
-                    _tx.rollback()
-                    continue
-                except:
-                    raise
+        finally:
+            setattr(_tx, "_exec_function_depth", depth)
+            # or if depth was 0, you might delete the attribute:
+            # if depth == 0:
+            #     delattr(_tx, "_exec_function_depth")
 
     @property
     def driver(self):
@@ -203,24 +179,36 @@ class Engine(object):
 
     @property
     def version(self):
+        """
+        Returns the DB server version.
+        """
         with Transaction(self) as tx:
             sql, vals = self.sql.version()
             return tx.execute(sql, vals).scalar()
 
     @property
     def timestamp(self):
+        """
+        Returns the current timestamp from the DB server.
+        """
         with Transaction(self) as tx:
             sql, vals = self.sql.timestamp()
             return tx.execute(sql, vals).scalar()
 
     @property
     def user(self):
+        """
+        Returns the current user as known by the DB server.
+        """
         with Transaction(self) as tx:
             sql, vals = self.sql.user()
             return tx.execute(sql, vals).scalar()
 
     @property
     def databases(self):
+        """
+        Returns a list of available databases.
+        """
         with Transaction(self) as tx:
             sql, vals = self.sql.databases()
             result = tx.execute(sql, vals)
@@ -228,13 +216,19 @@ class Engine(object):
 
     @property
     def current_database(self):
+        """
+        Returns the name of the current database.
+        """
         with Transaction(self) as tx:
             sql, vals = self.sql.current_database()
             return tx.execute(sql, vals).scalar()
 
     def create_database(self, name=None):
+        """
+        Creates a database if it doesn't exist, or does nothing if it does.
+        """
         old = None
-        if name == None:
+        if name is None:
             old = self.config["database"]
             self.set_config({"database": "postgres"})
             name = old
@@ -246,19 +240,27 @@ class Engine(object):
         return self
 
     def switch_to_database(self, database):
+        """
+        Switch the config to use a different database name, closing any existing connection.
+        """
         conf = self.config
         if "database" in conf:
             conf["database"] = database
         if "dbname" in conf:
             conf["dbname"] = database
-
         return self
 
     def set_config(self, config):
+        """
+        Updates the internal config dictionary.
+        """
         self.config.update(config)
 
     @property
     def schemas(self):
+        """
+        Returns a list of schemas in the current database.
+        """
         with Transaction(self) as tx:
             sql, vals = self.sql.schemas()
             result = tx.execute(sql, vals)
@@ -266,107 +268,106 @@ class Engine(object):
 
     @property
     def current_schema(self):
+        """
+        Returns the current schema in use.
+        """
         with Transaction(self) as tx:
             sql, vals = self.sql.current_schema()
             return tx.execute(sql, vals).scalar()
 
     @property
     def tables(self):
+        """
+        Returns a list of 'schema.table' for all tables in the current DB.
+        """
         with Transaction(self) as tx:
             sql, vals = self.sql.tables()
             result = tx.execute(sql, vals)
-            return ["%s.%s" % x for x in result.as_tuple()]
+            return [f"{x[0]}.{x[1]}" for x in result.as_tuple()]
 
     @property
     def views(self):
+        """
+        Returns a list of 'schema.view' for all views in the current DB.
+        """
         with Transaction(self) as tx:
             sql, vals = self.sql.views()
             result = tx.execute(sql, vals)
-            return ["%s.%s" % x for x in result.as_tuple()]
+            return [f"{x[0]}.{x[1]}" for x in result.as_tuple()]
 
-    def ProcessError(self, sql_stmt=None, sql_params=None):
-        sql = self.sql
+    def process_error(self, sql_stmt=None, sql_params=None):
+        """
+        Central method to parse driver exceptions and re-raise them as our custom exceptions.
+        """
         e = sys.exc_info()[1]
         msg = str(e).strip().lower()
+
         if isinstance(e, exceptions.DbException):
             raise
-        if hasattr(e, "pgcode"):
-            error_code = e.pgcode
-            error_mesg = e.pgerror
-        elif (
-            hasattr(e, "args") and isinstance(e.args, (tuple, list)) and len(e.args) > 1
-        ):
-            error_code = e[0]
-            error_mesg = e[1]
-        elif hasattr(e, "number") and hasattr(e, "text"):
-            error_code = e.number
-            error_mesg = e.text
-        elif hasattr(e, "args") and hasattr(e, "message"):
-            # SQLite3
-            error_code = None
-            error_mesg = e.message
-        else:
-            raise
-        if error_code in sql.ApplicationErrorCodes:
+
+        error_code, error_mesg = self.sql.get_error(e)
+
+        if error_code in self.sql.ApplicationErrorCodes:
             raise exceptions.DbApplicationError(e)
-        elif error_code in sql.ColumnMissingErrorCodes:
+        if error_code in self.sql.ColumnMissingErrorCodes:
             raise exceptions.DbColumnMissingError(e)
-        elif error_code in sql.TableMissingErrorCodes:
+        if error_code in self.sql.TableMissingErrorCodes:
             raise exceptions.DbTableMissingError(e)
-        elif error_code in sql.DatabaseMissingErrorCodes:
+        if error_code in self.sql.DatabaseMissingErrorCodes:
             raise exceptions.DbDatabaseMissingError(e)
-        elif error_code in sql.ForeignKeyMissingErrorCodes:
+        if error_code in self.sql.ForeignKeyMissingErrorCodes:
             raise exceptions.DbForeignKeyMissingError(e)
-        elif error_code in sql.TruncationErrorCodes:
+        if error_code in self.sql.TruncationErrorCodes:
             raise exceptions.DbTruncationError(e)
-        elif error_code in sql.DataIntegrityErrorCodes:
+        if error_code in self.sql.DataIntegrityErrorCodes:
             raise exceptions.DbDataIntegrityError(e)
-        elif error_code in sql.ConnectionErrorCodes:
+        if error_code in self.sql.ConnectionErrorCodes:
             raise exceptions.DbConnectionError(e)
-        elif error_code in sql.DuplicateKeyErrorCodes:
+        if error_code in self.sql.DuplicateKeyErrorCodes:
             raise exceptions.DbDuplicateKeyError(e)
-        elif re.search("key \(sys_id\)=\(\d+\) already exists.", msg, re.M):
+        if re.search(r"key \(sys_id\)=\(\d+\) already exists.", msg, re.M):
             raise exceptions.DbDuplicateKeyError(e)
-        elif error_code in sql.DatabaseObjectExistsErrorCodes:
+        if error_code in self.sql.DatabaseObjectExistsErrorCodes:
             raise exceptions.DbObjectExistsError(e)
-        elif error_code in sql.LockTimeoutErrorCodes:
+        if error_code in self.sql.LockTimeoutErrorCodes:
             raise exceptions.DbLockTimeoutError(e)
-        elif error_code in sql.RetryTransactionCodes:
+        if error_code in self.sql.RetryTransactionCodes:
             raise exceptions.DbRetryTransaction(e)
-        elif re.findall("database.*does not exist", msg, re.M):
+        if re.findall(r"database.*does not exist", msg, re.M):
             raise exceptions.DbDatabaseMissingError(e)
-        elif re.findall("no such database", msg, re.M):
+        if re.findall(r"no such database", msg, re.M):
             raise exceptions.DbDatabaseMissingError(e)
-        elif re.findall("already exists", msg, re.M):
+        if re.findall(r"already exists", msg, re.M):
             raise exceptions.DbObjectExistsError(e)
-        elif re.findall("server closed the connection unexpectedly", msg, re.M):
+        if re.findall(r"server closed the connection unexpectedly", msg, re.M):
             raise exceptions.DbConnectionError(e)
-        elif re.findall("no connection to the server", msg, re.M):
+        if re.findall(r"no connection to the server", msg, re.M):
             raise exceptions.DbConnectionError(e)
-        elif re.findall("connection timed out", msg, re.M):
+        if re.findall(r"connection timed out", msg, re.M):
             raise exceptions.DbConnectionError(e)
-        elif re.findall("could not connect to server", msg, re.M):
+        if re.findall(r"could not connect to server", msg, re.M):
             raise exceptions.DbConnectionError(e)
-        elif re.findall("cannot connect to server", msg, re.M):
+        if re.findall(r"cannot connect to server", msg, re.M):
             raise exceptions.DbConnectionError(e)
-        elif re.findall("connection already closed", msg, re.M):
+        if re.findall(r"connection already closed", msg, re.M):
             raise exceptions.DbConnectionError(e)
-        elif re.findall("cursor already closed", msg, re.M):
+        if re.findall(r"cursor already closed", msg, re.M):
             raise exceptions.DbConnectionError(e)
-        # SQLite3 errors
-        elif "no such table:" in msg:
+        if "no such table:" in msg:
             raise exceptions.DbTableMissingError(e)
-        print("Unhandled/Unknown Error in connection.ProcessError")
-        print("EXC_TYPE = {}".format(type(e)))
-        print("EXC_MSG = {}".format(str(e).strip()))
-        print("ERROR_CODE = {}".format(error_code))
-        print("ERROR_MSG = {}".format(error_mesg))
-        if sql_stmt:
-            print("\n")
-            print("sql_stmt [velocity.db.engine]: {}".format(sql_stmt))
-            print("\n")
-        if sql_params:
-            print(sql_params)
-            print("\n")
-        traceback.print_exc()
+
+        msg = f"""
+Unhandled/Unknown Error in engine.process_error
+EXC_TYPE = {type(e)}
+EXC_MSG = {str(e).strip()}
+
+ERROR_CODE = {error_code}
+ERROR_MSG = {error_mesg}
+
+SQL_STMT = {sql_stmt}
+SQL_PARAMS = {sql_params}
+
+{traceback.format_exc()}
+"""
+        print(msg)
         raise