sqla-fancy-core 1.0.0__py3-none-any.whl → 1.2.2__py3-none-any.whl

sqla_fancy_core/__init__.py

@@ -1,159 +1,11 @@
 """SQLAlchemy core, but fancier."""
 
-from typing import Optional, Union, overload
-
-import sqlalchemy as sa
-
-
-class TableFactory:
-    """A factory for creating SQLAlchemy columns with default values."""
-
-    def __init__(self, metadata: Optional[sa.MetaData] = None):
-        """Initialize the factory with default values."""
-        if metadata is None:
-            metadata = sa.MetaData()
-        self.metadata = metadata
-        self.c = []
-
-    def col(self, *args, **kwargs) -> sa.Column:
-        col = sa.Column(*args, **kwargs)
-        self.c.append(col)
-        return col
-
-    def integer(self, name: str, *args, **kwargs) -> sa.Column:
-        return self.col(name, sa.Integer, *args, **kwargs)
-
-    def string(self, name: str, *args, **kwargs) -> sa.Column:
-        return self.col(name, sa.String, *args, **kwargs)
-
-    def text(self, name: str, *args, **kwargs) -> sa.Column:
-        return self.col(name, sa.Text, *args, **kwargs)
-
-    def float(self, name: str, *args, **kwargs) -> sa.Column:
-        return self.col(name, sa.Float, *args, **kwargs)
-
-    def numeric(self, name: str, *args, **kwargs) -> sa.Column:
-        return self.col(name, sa.Numeric, *args, **kwargs)
-
-    def bigint(self, name: str, *args, **kwargs) -> sa.Column:
-        return self.col(name, sa.BigInteger, *args, **kwargs)
-
-    def smallint(self, name: str, *args, **kwargs) -> sa.Column:
-        return self.col(name, sa.SmallInteger, *args, **kwargs)
-
-    def timestamp(self, name: str, *args, **kwargs) -> sa.Column:
-        return self.col(name, sa.TIMESTAMP, *args, **kwargs)
-
-    def date(self, name: str, *args, **kwargs) -> sa.Column:
-        return self.col(name, sa.Date, *args, **kwargs)
-
-    def datetime(self, name: str, *args, **kwargs) -> sa.Column:
-        return self.col(name, sa.DateTime, *args, **kwargs)
-
-    def today(self, name: str, *args, **kwargs) -> sa.Column:
-        return self.date(name, default=sa.func.now(), *args, **kwargs)
-
-    def time(self, name: str, *args, **kwargs) -> sa.Column:
-        return self.col(name, sa.Time, *args, **kwargs)
-
-    def timenow(self, name: str, *args, **kwargs) -> sa.Column:
-        return self.time(name, default=sa.func.now(), *args, **kwargs)
-
-    def now(self, name: str, *args, **kwargs) -> sa.Column:
-        return self.datetime(name, default=sa.func.now(), *args, **kwargs)
-
-    def boolean(self, name: str, *args, **kwargs) -> sa.Column:
-        return self.col(name, sa.Boolean, *args, **kwargs)
-
-    def true(self, name: str, *args, **kwargs):
-        return self.boolean(name, default=True, *args, **kwargs)
-
-    def false(self, name: str, *args, **kwargs):
-        return self.boolean(name, default=False, *args, **kwargs)
-
-    def foreign_key(self, name: str, ref: Union[str, sa.Column], *args, **kwargs):
-        return self.col(name, sa.ForeignKey(ref), *args, **kwargs)
-
-    def enum(self, name: str, enum: type, *args, **kwargs) -> sa.Column:
-        return self.col(name, sa.Enum(enum), *args, **kwargs)
-
-    def json(self, name: str, *args, **kwargs) -> sa.Column:
-        return self.col(name, sa.JSON, *args, **kwargs)
-
-    def array(self, name: str, *args, **kwargs) -> sa.Column:
-        return self.col(name, sa.ARRAY, *args, **kwargs)
-
-    def array_int(self, name: str, *args, **kwargs) -> sa.Column:
-        return self.array(name, sa.Integer, *args, **kwargs)
-
-    def array_str(self, name: str, *args, **kwargs) -> sa.Column:
-        return self.array(name, sa.String, *args, **kwargs)
-
-    def array_text(self, name: str, *args, **kwargs) -> sa.Column:
-        return self.array(name, sa.Text, *args, **kwargs)
-
-    def array_float(self, name: str, *args, **kwargs) -> sa.Column:
-        return self.array(name, sa.Float, *args, **kwargs)
-
-    def array_numeric(self, name: str, *args, **kwargs) -> sa.Column:
-        return self.array(name, sa.Numeric, *args, **kwargs)
-
-    def array_bigint(self, name: str, *args, **kwargs) -> sa.Column:
-        return self.array(name, sa.BigInteger, *args, **kwargs)
-
-    def array_smallint(self, name: str, *args, **kwargs) -> sa.Column:
-        return self.array(name, sa.SmallInteger, *args, **kwargs)
-
-    def array_timestamp(self, name: str, *args, **kwargs) -> sa.Column:
-        return self.array(name, sa.TIMESTAMP, *args, **kwargs)
-
-    def array_date(self, name: str, *args, **kwargs) -> sa.Column:
-        return self.array(name, sa.Date, *args, **kwargs)
-
-    def array_datetime(self, name: str, *args, **kwargs) -> sa.Column:
-        return self.array(name, sa.DateTime, *args, **kwargs)
-
-    def array_time(self, name: str, *args, **kwargs) -> sa.Column:
-        return self.array(name, sa.Time, *args, **kwargs)
-
-    def array_boolean(self, name: str, *args, **kwargs) -> sa.Column:
-        return self.array(name, sa.Boolean, *args, **kwargs)
-
-    def array_enum(self, name: str, enum: type, *args, **kwargs) -> sa.Column:
-        return self.array(name, sa.Enum(enum), *args, **kwargs)
-
-    def auto_id(self, name="id", *args, **kwargs) -> sa.Column:
-        return self.integer(
-            name, primary_key=True, index=True, autoincrement=True, *args, **kwargs
-        )
-
-    def updated_at(self, name="updated_at", *args, **kwargs) -> sa.Column:
-        return self.datetime(
-            name, default=sa.func.now(), onupdate=sa.func.now(), *args, **kwargs
-        )
-
-    def created_at(self, name="created_at", *args, **kwargs) -> sa.Column:
-        return self.datetime(name, default=sa.func.now(), *args, **kwargs)
-
-    @overload
-    def __call__(self, arg1: str, *args, **kwargs) -> sa.Table: ...
-    @overload
-    def __call__(self, arg1: sa.Column, *args, **kwargs) -> sa.Column: ...
-    @overload
-    def __call__(self, arg1: sa.Table, *args, **kwargs) -> sa.Table: ...
-    def __call__(self, arg1, *args, **kwargs):
-        if isinstance(arg1, str):
-            cols = self.c
-            self.c = []
-            return sa.Table(arg1, self.metadata, *args, *cols, **kwargs)
-        elif isinstance(arg1, sa.Column):
-            arg1.info["args"] = args
-            arg1.info["kwargs"] = kwargs
-            self.c.append(arg1)
-            return arg1
-        elif isinstance(arg1, sa.Table):
-            cols = self.c
-            self.c = []
-            return sa.Table(arg1.name, self.metadata, *args, *cols, **kwargs)
-        else:
-            raise TypeError(f"Expected a string or Column, got {type(arg1).__name__}")
+from sqla_fancy_core.factories import TableFactory  # noqa
+from sqla_fancy_core.wrappers import (  # noqa
+    FancyEngineWrapper,
+    AsyncFancyEngineWrapper,
+    fancy,
+    FancyError,
+    AtomicContextError,
+)
+from sqla_fancy_core.decorators import transact, Inject  # noqa

sqla_fancy_core/decorators.py

@@ -0,0 +1,196 @@
+"""Some decorators for fun times with SQLAlchemy core."""
+
+import functools
+import inspect
+from typing import Union, overload
+
+import sqlalchemy as sa
+from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine
+
+from sqla_fancy_core.wrappers import AsyncFancyEngineWrapper, FancyEngineWrapper
+
+EngineType = Union[sa.Engine, AsyncEngine, FancyEngineWrapper, AsyncFancyEngineWrapper]
+
+
+class _Injectable:
+    def __init__(self, engine: EngineType):
+        self.engine = engine
+
+
+@overload
+def Inject(engine: Union[sa.Engine, FancyEngineWrapper]) -> sa.Connection: ...
+@overload
+def Inject(engine: Union[AsyncEngine, AsyncFancyEngineWrapper]) -> AsyncConnection: ...
+def Inject(engine: EngineType):  # type: ignore
+    """A marker class for dependency injection."""
+    return _Injectable(engine)
+
+
+def transact(func):
+    """A decorator that provides a transactional context.
+
+    If the decorated function is called with a connection object, that
+    connection is used. Otherwise, a new transaction is started from the
+    engine, and the new connection is injected to the function.
+
+    Example: ::
+        @transact
+        def create_user(name: str, conn: sa.Connection = Inject(engine)):
+            conn.execute(...)
+
+        # This will create a new transaction
+        create_user("test")
+
+        # This will use the existing connection
+        with engine.connect() as conn:
+            create_user(name="existing", conn=conn)
+    """
+
+    sig = inspect.signature(func)
+    inject_param_name = None
+    for name, param in sig.parameters.items():
+        if param.default is not inspect.Parameter.empty and isinstance(
+            param.default, _Injectable
+        ):
+            inject_param_name = name
+            break
+    if inject_param_name is None:
+        return func  # No injection needed
+
+    engine_arg = sig.parameters[inject_param_name].default.engine
+    if isinstance(engine_arg, sa.Engine):
+        is_async = False
+        engine = engine_arg
+    elif isinstance(engine_arg, FancyEngineWrapper):
+        is_async = False
+        engine = engine_arg.engine
+    elif isinstance(engine_arg, AsyncEngine):
+        is_async = True
+        engine = engine_arg
+    elif isinstance(engine_arg, AsyncFancyEngineWrapper):
+        is_async = True
+        engine = engine_arg.engine
+    else:
+        raise TypeError("Unsupported engine type")
+
+    if is_async:
+
+        @functools.wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            conn = kwargs.get(inject_param_name)
+            if isinstance(conn, AsyncConnection):
+                if conn.in_transaction():
+                    return await func(*args, **kwargs)
+                else:
+                    async with conn.begin():
+                        return await func(*args, **kwargs)
+            elif isinstance(conn, sa.Connection):
+                if conn.in_transaction():
+                    return await func(*args, **kwargs)
+                else:
+                    with conn.begin():
+                        return await func(*args, **kwargs)
+            else:
+                async with engine.begin() as conn:  # type: ignore
+                    kwargs[inject_param_name] = conn
+                    return await func(*args, **kwargs)
+
+        return async_wrapper
+
+    else:
+
+        @functools.wraps(func)
+        def sync_wrapper(*args, **kwargs):
+            conn = kwargs.get(inject_param_name)
+            if isinstance(conn, sa.Connection):
+                if conn.in_transaction():
+                    return func(*args, **kwargs)
+                else:
+                    with conn.begin():
+                        return func(*args, **kwargs)
+            elif isinstance(conn, AsyncConnection):
+                raise TypeError("AsyncConnection cannot be used in sync function")
+            else:
+                with engine.begin() as conn:  # type: ignore
+                    kwargs[inject_param_name] = conn
+                    return func(*args, **kwargs)
+
+        return sync_wrapper
+
+
+def connect(func):
+    """A decorator that provides a connection context.
+
+    If the decorated function is called with a connection object, that
+    connection is used. Otherwise, a new connection is created from the
+    engine, and the new connection is injected to the function.
+
+    Example: ::
+        @connect
+        def get_user_count(conn: sa.Connection = Inject(engine)) -> int:
+            return conn.execute(...).scalar_one()
+
+        # This will create a new connection
+        count = get_user_count()
+
+        # This will use the existing connection
+        with engine.connect() as conn:
+            count = get_user_count(conn)
+    """
+
+    sig = inspect.signature(func)
+    inject_param_name = None
+    for name, param in sig.parameters.items():
+        if param.default is not inspect.Parameter.empty and isinstance(
+            param.default, _Injectable
+        ):
+            inject_param_name = name
+            break
+    if inject_param_name is None:
+        return func  # No injection needed
+
+    engine_arg = sig.parameters[inject_param_name].default.engine
+    if isinstance(engine_arg, sa.Engine):
+        is_async = False
+        engine = engine_arg
+    elif isinstance(engine_arg, FancyEngineWrapper):
+        is_async = False
+        engine = engine_arg.engine
+    elif isinstance(engine_arg, AsyncEngine):
+        is_async = True
+        engine = engine_arg
+    elif isinstance(engine_arg, AsyncFancyEngineWrapper):
+        is_async = True
+        engine = engine_arg.engine
+    else:
+        raise TypeError("Unsupported engine type")
+
+    if is_async:
+
+        @functools.wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            conn = kwargs.get(inject_param_name)
+            if isinstance(conn, (AsyncConnection, sa.Connection)):
+                return await func(*args, **kwargs)
+            else:
+                async with engine.connect() as conn:  # type: ignore
+                    kwargs[inject_param_name] = conn
+                    return await func(*args, **kwargs)
+
+        return async_wrapper
+
+    else:
+
+        @functools.wraps(func)
+        def sync_wrapper(*args, **kwargs):
+            conn = kwargs.get(inject_param_name)
+            if isinstance(conn, sa.Connection):
+                return func(*args, **kwargs)
+            elif isinstance(conn, AsyncConnection):
+                raise TypeError("AsyncConnection cannot be used in sync function")
+            else:
+                with engine.connect() as conn:  # type: ignore
+                    kwargs[inject_param_name] = conn
+                    return func(*args, **kwargs)
+
+        return sync_wrapper

sqla_fancy_core/factories.py

@@ -0,0 +1,156 @@
+"""Some factories for fun times with SQLAlchemy core."""
+
+from typing import Optional, Union, overload
+
+import sqlalchemy as sa
+
+
+class TableFactory:
+    """A factory for creating SQLAlchemy columns with default values."""
+
+    def __init__(self, metadata: Optional[sa.MetaData] = None):
+        """Initialize the factory with default values."""
+        if metadata is None:
+            self.metadata = sa.MetaData()
+        else:
+            self.metadata = metadata
+        self.c = []
+
+    def col(self, *args, **kwargs) -> sa.Column:
+        col = sa.Column(*args, **kwargs)
+        return self(col)
+
+    def integer(self, name: str, *args, **kwargs) -> sa.Column:
+        return self.col(name, sa.Integer, *args, **kwargs)
+
+    def string(self, name: str, *args, **kwargs) -> sa.Column:
+        return self.col(name, sa.String, *args, **kwargs)
+
+    def text(self, name: str, *args, **kwargs) -> sa.Column:
+        return self.col(name, sa.Text, *args, **kwargs)
+
+    def float(self, name: str, *args, **kwargs) -> sa.Column:
+        return self.col(name, sa.Float, *args, **kwargs)
+
+    def numeric(self, name: str, *args, **kwargs) -> sa.Column:
+        return self.col(name, sa.Numeric, *args, **kwargs)
+
+    def bigint(self, name: str, *args, **kwargs) -> sa.Column:
+        return self.col(name, sa.BigInteger, *args, **kwargs)
+
+    def smallint(self, name: str, *args, **kwargs) -> sa.Column:
+        return self.col(name, sa.SmallInteger, *args, **kwargs)
+
+    def timestamp(self, name: str, *args, **kwargs) -> sa.Column:
+        return self.col(name, sa.TIMESTAMP, *args, **kwargs)
+
+    def date(self, name: str, *args, **kwargs) -> sa.Column:
+        return self.col(name, sa.Date, *args, **kwargs)
+
+    def datetime(self, name: str, *args, **kwargs) -> sa.Column:
+        return self.col(name, sa.DateTime, *args, **kwargs)
+
+    def today(self, name: str, *args, **kwargs) -> sa.Column:
+        return self.date(name, default=sa.func.now(), *args, **kwargs)
+
+    def time(self, name: str, *args, **kwargs) -> sa.Column:
+        return self.col(name, sa.Time, *args, **kwargs)
+
+    def timenow(self, name: str, *args, **kwargs) -> sa.Column:
+        return self.time(name, default=sa.func.now(), *args, **kwargs)
+
+    def now(self, name: str, *args, **kwargs) -> sa.Column:
+        return self.datetime(name, default=sa.func.now(), *args, **kwargs)
+
+    def boolean(self, name: str, *args, **kwargs) -> sa.Column:
+        return self.col(name, sa.Boolean, *args, **kwargs)
+
+    def true(self, name: str, *args, **kwargs):
+        return self.boolean(name, default=True, *args, **kwargs)
+
+    def false(self, name: str, *args, **kwargs):
+        return self.boolean(name, default=False, *args, **kwargs)
+
+    def foreign_key(self, name: str, ref: Union[str, sa.Column], *args, **kwargs):
+        return self.col(name, sa.ForeignKey(ref), *args, **kwargs)
+
+    def enum(self, name: str, enum: type, *args, **kwargs) -> sa.Column:
+        return self.col(name, sa.Enum(enum), *args, **kwargs)
+
+    def json(self, name: str, *args, **kwargs) -> sa.Column:
+        return self.col(name, sa.JSON, *args, **kwargs)
+
+    def array(self, name: str, *args, **kwargs) -> sa.Column:
+        return self.col(name, sa.ARRAY, *args, **kwargs)
+
+    def array_int(self, name: str, *args, **kwargs) -> sa.Column:
+        return self.array(name, sa.Integer, *args, **kwargs)
+
+    def array_str(self, name: str, *args, **kwargs) -> sa.Column:
+        return self.array(name, sa.String, *args, **kwargs)
+
+    def array_text(self, name: str, *args, **kwargs) -> sa.Column:
+        return self.array(name, sa.Text, *args, **kwargs)
+
+    def array_float(self, name: str, *args, **kwargs) -> sa.Column:
+        return self.array(name, sa.Float, *args, **kwargs)
+
+    def array_numeric(self, name: str, *args, **kwargs) -> sa.Column:
+        return self.array(name, sa.Numeric, *args, **kwargs)
+
+    def array_bigint(self, name: str, *args, **kwargs) -> sa.Column:
+        return self.array(name, sa.BigInteger, *args, **kwargs)
+
+    def array_smallint(self, name: str, *args, **kwargs) -> sa.Column:
+        return self.array(name, sa.SmallInteger, *args, **kwargs)
+
+    def array_timestamp(self, name: str, *args, **kwargs) -> sa.Column:
+        return self.array(name, sa.TIMESTAMP, *args, **kwargs)
+
+    def array_date(self, name: str, *args, **kwargs) -> sa.Column:
+        return self.array(name, sa.Date, *args, **kwargs)
+
+    def array_datetime(self, name: str, *args, **kwargs) -> sa.Column:
+        return self.array(name, sa.DateTime, *args, **kwargs)
+
+    def array_time(self, name: str, *args, **kwargs) -> sa.Column:
+        return self.array(name, sa.Time, *args, **kwargs)
+
+    def array_boolean(self, name: str, *args, **kwargs) -> sa.Column:
+        return self.array(name, sa.Boolean, *args, **kwargs)
+
+    def array_enum(self, name: str, enum: type, *args, **kwargs) -> sa.Column:
+        return self.array(name, sa.Enum(enum), *args, **kwargs)
+
+    def auto_id(self, name="id", *args, **kwargs) -> sa.Column:
+        return self.integer(
+            name, primary_key=True, index=True, autoincrement=True, *args, **kwargs
+        )
+
+    def updated_at(self, name="updated_at", *args, **kwargs) -> sa.Column:
+        return self.datetime(
+            name, default=sa.func.now(), onupdate=sa.func.now(), *args, **kwargs
+        )
+
+    def created_at(self, name="created_at", *args, **kwargs) -> sa.Column:
+        return self.datetime(name, default=sa.func.now(), *args, **kwargs)
+
+    @overload
+    def __call__(self, arg1: str, *args, **kwargs) -> sa.Table: ...
+    @overload
+    def __call__(self, arg1: sa.Column, *args, **kwargs) -> sa.Column: ...
+    @overload
+    def __call__(self, arg1: sa.Table, *args, **kwargs) -> sa.Table: ...
+    def __call__(self, arg1, *args, **kwargs):
+        if isinstance(arg1, sa.Column):
+            arg1.info["args"] = args
+            arg1.info["kwargs"] = kwargs
+            self.c.append(arg1)
+            return arg1
+        elif isinstance(arg1, Union[str, sa.Table]):
+            cols = self.c
+            tablename = arg1.name if isinstance(arg1, sa.Table) else arg1
+            self.c = []
+            return sa.Table(tablename, self.metadata, *args, *cols, **kwargs)
+        else:
+            raise TypeError(f"Expected a string or Column, got {type(arg1).__name__}")