New QueryBuilder and EngineDispatcher classes. Improved create_initial_migration()

Author: maximebf

docs/engine.md

@@ -57,3 +57,33 @@ Connections are pooled and re-used by default. You can disabled this behavior by
 `max_pool_conns` can also be used to define the maximum number of connections to start.
 
 Use `engine.disconnect_all()` to close all connections.
+
+## Engine dispatcher
+
+Multiple engines can be used at the same time. `EngineDispatcher` makes it easy to select engines based on matching tags.
+
+Register engines on the dispatcher using `dispatcher.register(engine, ["tag1", "tag2"])`.
+
+Select one engine using `dispatcher.tag_name`. If multiple engines matche a tag, one is randomly selected.
+
+```py
+from sqlorm import EngineDispatcher
+
+dispatcher = EngineDispatcher()
+dispatcher.register(Engine.from_uri("postgresql://primary"), default=True) # there can be multiple default engines
+dispatcher.register(Engine.from_uri("postgresql://replica1"), ["readonly"])
+dispatcher.register(Engine.from_uri("postgresql://replica2"), ["readonly"])
+
+with dispatcher:
+    # uses default engine (ie. primary)
+
+with dispatcher.readonly:
+    # uses a randomly selected engines from the one matching the readonly tag (ie. replica1 or replica2)
+```
+
+The context behaves the same as a context from an Engine, ie. starting a transaction. To start a session instead, use `dispatcher.session(tag_name)`.
+
+By default, if no engines match the given tag, it will fallback on the default. This can be changed using the `fallback` argument of the `EngineDispatcher` constructor.
+Pass `False` to disable fallback or a tag name to fallback to a specific tag.
+
+Use `dispatcher.disconnect_all()` to close all connections across all engines.

docs/executing.md

@@ -61,7 +61,7 @@ There are a few shortcut methods on the transaction:
 ```python
 with engine as tx:
     task_row = tx.fetch("SELECT * FROM tasks WHERE id = ?", [1]).first()
-    task_row = tx.fetchone("SELECT * FROM tasks WHERE id = ?", [1]) # same as line before
+    task_row = tx.fetchone("SELECT * FROM tasks WHERE id = ?", [1]) # same as previous line
 
     for row in tx.fetch("SELECT * FROM tasks"):
         print(row)

docs/integrations.md

@@ -0,0 +1,5 @@
+# Integrations
+
+## Flask
+
+See [Flask-SQLORM](https://github.com/hyperflask/flask-sqlorm).

docs/sql-utilities.md

@@ -74,6 +74,21 @@ Useful shortcuts:
 - `SQL.update()` instead of `SQL().update()`
 - `SQL.delete_from()` instead of `SQL().delete_from()`
 
+## Query builder
+
+A query builder for SELECT statements, built on top of the `SQL` class, is provided.
+
+```py
+from sqlorm import QueryBuilder
+query = QueryBuilder().from_("table").where(col="value").where(SQL("col") >= 1)
+```
+
+It is aware of the different parts of the query and will treat them differently.
+
+ - `builder.select()`, `.join()`, `.where()` will add expressions, all other parts will be replaced when called
+ - If `.select()` is not called, `SELECT *` is used
+ - Keyword arguments can be used with `.where()` to generate equal comparisons
+
 ## Handling list of SQL pieces
 
 Use `SQL.List` to manage lists of `SQL` objects. A few subclasses exists:

mkdocs.yml

@@ -16,6 +16,7 @@ nav:
   - schema.md
   - drivers.md
   - instrumentation.md
+  - integrations.md
 
 theme:
   name: material

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "sqlorm-py"
-version = "0.3.1"
+version = "0.4.0"
 description = "A new kind or ORM that do not abstract away your database or SQL queries."
 authors = [
     {"name" = "Maxime Bouroumeau-Fuseau", email = "[email protected]"}

src/sqlorm/__init__.py

@@ -1,6 +1,7 @@
 from .engine import (
     Engine,
     EngineError,
+    EngineDispatcher,
     Session,
     SessionError,
     Transaction,

src/sqlorm/builder.py

@@ -11,11 +11,12 @@ class QueryBuilderError(Exception):
 
 
 class QueryBuilder(SQLStr):
-    components = ("SELECT", "FROM", "JOIN+", "WHERE+", "GROUP BY", "HAVING", "ORDER BY", "LIMIT", "OFFSET")
+    components = ("SELECT+", "FROM", "JOIN+", "WHERE+", "GROUP BY", "HAVING", "ORDER BY", "LIMIT", "OFFSET")
     wrappers = {"SELECT": List, "FROM": List, "WHERE": Conditions, "HAVING": Conditions, "ORDER BY": List}
 
     def __init__(self):
-        self.parts = {}
+        self.parts = {"SELECT": QueryPartBuilder(self, "*")}
+        self._selected = False
 
     def _render(self, params):
         stmt = []
@@ -38,12 +39,23 @@ def __getattr__(self, name):
 
         self.parts[name] = QueryPartBuilder(self)
         return self.parts[name]
+    
+    def select(self, *parts):
+        if not self._selected:
+            self._selected = True
+            self.parts["SELECT"] = QueryPartBuilder(self)
+        return self.parts["SELECT"](*parts)
+    
+    def where(self, *parts, **filters):
+        for k, v in filters.items():
+            parts = list(parts) + [SQL(k) == v]
+        return self.__getattr__("where")(*parts)
 
 
 class QueryPartBuilder:
-    def __init__(self, builder):
+    def __init__(self, builder, *parts):
         self.builder = builder
-        self.parts = []
+        self.parts = list(parts)
 
     def __call__(self, *parts):
         if len(parts) == 1 and parts[0] is None:

src/sqlorm/engine.py

@@ -4,6 +4,7 @@
 import logging
 import inspect
 import urllib.parse
+import random
 from blinker import Namespace
 from .sql import render, ParametrizedStmt
 from .resultset import ResultSet, CompositeResultSet, CompositionMap
@@ -152,16 +153,68 @@ def session(self, **kwargs):
             session.close()
 
     def __enter__(self):
-        if session_context.top:
+        if session_context.top and session_context.top.engine == self:
             return session_context.top.__enter__()
-        session = self.make_session()
+        session = self.make_session(close_after_tx=True)
         return session.__enter__()
 
     def __exit__(self, exc_type, exc_value, exc_tb):
-        session = session_context.top
-        session.__exit__(exc_type, exc_value, exc_tb)
-        if not session.transactions:
-            session.close()
+        session_context.top.__exit__(exc_type, exc_value, exc_tb)
+
+
+class EngineDispatcher:
+    def __init__(self, fallback=None, default_tag='default'):
+        self._engines = []
+        self.default_tag = default_tag
+        self.fallback = default_tag if fallback is None else fallback
+
+    def register(self, engine, tags=None, default=False):
+        tags = tags or []
+        if default:
+            tags.append(self.default_tag)
+        self._engines.append((engine, tags))
+
+    @property
+    def engines(self):
+        return [e for e, tags in self.engines]
+
+    @property
+    def tags(self):
+        return list(set([t for e, tags in self.engines for t in tags]))
+
+    def select_all(self, tag=None):
+        if not tag and self.default_tag:
+            tag = self.default_tag
+        found = []
+        for engine, tags in self._engines:
+            if not tag or tag in tags:
+                found.append(engine)
+        if found:
+            return found
+        if self.fallback is True:
+            return self.engines
+        if self.fallback:
+            return self.select_all(self.fallback)
+        raise EngineError("No engines found")
+
+    def select(self, tag=None):
+        return random.choice(self.select_all(tag))
+    
+    def session(self, tag=None, **kwargs):
+        return self.select(tag).session(**kwargs)
+    
+    def disconnect_all(self):
+        for engine in self.engines:
+            engine.disconnect_all()
+
+    def __getattr__(self, tag):
+        return self.select(tag)
+    
+    def __enter__(self):
+        return self.select().__enter__()
+
+    def __exit__(self, exc_type, exc_value, exc_tb):
+        session_context.top.__exit__(exc_type, exc_value, exc_tb)
 
 
 class EngineError(Exception):
@@ -258,6 +311,7 @@ def __init__(
         dbapi_conn=None,
         auto_close_conn=None,
         virtual_tx=False,
+        close_after_tx=False,
         logger=None,
         logger_level=None,
         engine=None,
@@ -274,6 +328,7 @@ def __init__(
         self.conn = dbapi_conn
         self.auto_close_conn = not bool(dbapi_conn) if auto_close_conn is None else auto_close_conn
         self.virtual_tx = virtual_tx
+        self.close_after_tx = close_after_tx
         self.logger = logger
         self.logger_level = logger_level
         self.transactions = []
@@ -344,6 +399,8 @@ def __exit__(self, exc_type, exc_value, exc_tb):
         tx.__exit__(exc_type, exc_value, exc_tb)
         if not self.transactions:
             session_context.pop()
+            if self.close_after_tx:
+                self.close()
 
     @property
     def transaction(self):

src/sqlorm/schema.py

@@ -43,7 +43,7 @@ def create_table(mapper, default_type="varchar"):
     return SQL(
         "CREATE TABLE",
         mapper.table,
-        SQL.Tuple(
+        SQL.List(
             [
                 SQL(c.schema_def) if c.schema_def else SQL(
                     c.name,
@@ -61,7 +61,10 @@ def create_table(mapper, default_type="varchar"):
                     else "",
                 )
                 for c in mapper.columns
-            ]
+            ],
+            startstr="(\n   ",
+            joinstr=",\n   ",
+            endstr="\n)"
         ),
     )
 
@@ -149,14 +152,29 @@ def set_schema_version(version, engine=None):
             )
 
 
-def create_initial_migration(model_registry=None, path="migrations", version="000", **kwargs):
+def create_initial_migration(model_registry=None, models=None, path="migrations", version=None, **kwargs):
     if not model_registry:
         model_registry = BaseModel.__model_registry__
 
+    if version is None:
+        migrations = create_migrations_from_dir(path)
+        version = '%03d' % (migrations[-1][0] + 1) if migrations else "000"
+
+    if not models:
+        models = model_registry.values()
+        header = "Initial creation of the database"
+        name = "initial"
+    else:
+        models = [model_registry[m] if isinstance(m, str) else m for m in models]
+        header = "Creation of tables: " + ", ".join(m.__mapper__.table for m in models)
+        name = "create_" + "_".join(m.__mapper__.table for m in models)
+
     stmts = []
-    for model in model_registry.values():
-        stmts.append(create_table.sql(model.__mapper__, **kwargs))
+    for model in models:
+        stmts.append(str(create_table.sql(model.__mapper__, **kwargs)))
+
+    sql = f"-- {header} (auto-generated by sqlorm)\n\n" + ";\n\n".join(stmts) + ";\n"
+    with open(os.path.join(path, f"{version}_{name}.sql"), "w") as f:
+        f.write(sql)
 
-    sql = "-- Initial creation of the database (auto-generated by sqlorm)\n\n" + ";\n\n".join(stmts)
-    with open(os.path.join(path, f"{version}_initial.sql"), "w") as f:
-        f.write(sql)
+    return (header, version, name)