Set INSERT and UPDATE table alias in SQLAlchemy

user_2340939 0
Viewed 126

How to set alias for the table in INSERT and UPDATE clauses in SQLAlchemy? Here is a PostgreSQL example that utilizes both.

CREATE TABLE person (
  id INT PRIMARY KEY,
  name CHAR(120)
);

INSERT INTO person (id, name) 
VALUES 
  (1, 'A'), 
  (2, 'B'), 
  (3, 'C'), 
  (4, 'D');

INSERT INTO person AS p (id, name) 
VALUES (1, 'A2')
ON CONFLICT (id) DO UPDATE 
SET name = EXCLUDED.name || ' (previously ' || p.name || ')';

UPDATE person AS p 
SET name = 'B2' || ' (previously ' || p.name || ')'
WHERE p.id = 2;

SELECT * 
FROM person 
ORDER BY id;
sqlalchemy postgresql orm python
1 Answers

SQLAlchemy's Insert and Update constructs expect Table* instances as their primary arguments and do not provide support for aliasing the names of these objects. Aliased objects created using the alias function do not expose all the attributes required for statement compilation; passing such an object to insert or update results in errors.

However aliased objects expose the table that they are aliasing, so we can hook into the compilation process and patch them before the compiler complains.

Update requires that we patch a single attribute:

@compiles(Update)
def visit_update_alias(update, compiler, **kw):
    if isinstance(update.table, Alias):
        update.table.implicit_returning = update.table.original.implicit_returning
    return compiler.visit_update(update, **kw)

Insert is more complicated. There are three attributes to patch, and we must post-process the compiled statement to insert the table name. Even with patched attribute it raises a warning about the primary key column†.

@compiles(Insert)
def visit_insert_into_alias(insert, compiler, **kw):
    if isinstance(insert.table, Alias):
        alias = insert.table
        table = insert.table.original
        alias._autoincrement_column = table._autoincrement_column
        alias.fullname = table.fullname
        alias.implicit_returning = table.implicit_returning
        prepared_table_name = compiler.preparer.format_table(table)
        compiled = compiler.visit_insert(insert, **kw)
        compiled = compiled.replace(
            f'INSERT INTO {alias.name}',
            f'INSERT INTO {prepared_table_name} AS {alias.name}',
        )
        return compiled
    return compiler.visit_insert(insert, **kw)

Obviously these are barely-tested hacks, and brittle in the face of changes to SQLAlchemy's internals or more complex use cases. If you really need this functionality I'd recommend opening a discussion on GitHub to see if it could be added to the library.

Here's a complete script for the examples in the question (plus a plain insert).

import sqlalchemy as sa
from sqlalchemy.dialects.postgresql import insert
from sqlalchemy.ext.compiler import compiles
from sqlalchemy.sql.expression import Insert, Update
from sqlalchemy.sql.selectable import Alias


@compiles(Insert)
def visit_insert_into_alias(insert, compiler, **kw):
    if isinstance(insert.table, Alias):
        alias = insert.table
        table = insert.table.original
        alias._autoincrement_column = table._autoincrement_column
        alias.fullname = table.fullname
        alias.implicit_returning = table.implicit_returning
        prepared_table_name = compiler.preparer.format_table(table)
        compiled = compiler.visit_insert(insert, **kw)
        compiled = compiled.replace(
            f'INSERT INTO {alias.name}',
            f'INSERT INTO {prepared_table_name} AS {alias.name}',
        )
        return compiled
    return compiler.visit_insert(insert, **kw)


@compiles(Update)
def visit_update_alias(update, compiler, **kw):
    if isinstance(update.table, Alias):
        update.table.implicit_returning = update.table.original.implicit_returning
    return compiler.visit_update(update, **kw)


tbl = sa.Table(
    'person',
    sa.MetaData(),
    sa.Column('id', sa.Integer, primary_key=True),
    sa.Column('name', sa.String),
)


engine = sa.create_engine(
    'postgresql+psycopg2:///test', echo=True, future=True
)
tbl.drop(engine, checkfirst=True)
tbl.create(engine)

with engine.begin() as conn:
    conn.execute(tbl.insert(), [{'name': c} for c in 'ABCD'])

with engine.begin() as conn:
    pal = sa.alias(tbl, name='p')
    conn.execute(insert(pal), [{'name': c} for c in 'EFGH'])

with engine.begin() as conn:
    p = sa.alias(tbl, name='p')
    ins = insert(p).values(id=1, name='A2')
    upd = ins.on_conflict_do_update(
        index_elements=['id'],
        set_=dict(name=ins.excluded.name + ' previously ' + p.c.name),
    )
    conn.execute(upd)

with engine.begin() as conn:
    p = sa.alias(tbl, name='p')
    upd = (
        sa.update(p)
        .where(p.c.id == 2)
        .values(name='B2 previously ' + p.c.name)
    )
    conn.execute(upd)

with engine.connect() as conn:
    rows = conn.execute(sa.select(tbl))
    for row in rows:
        print(row)

* Technically, Insert expects a TableClause, from which Table inherits.

† The warning is

Column 'person.id' is marked as a member of the primary key for table 'person', but has no Python-side or server-side default generator indicated, nor does it indicate 'autoincrement=True' or 'nullable=True', and no explicit value is passed. Primary key columns typically may not store NULL.

I think this is a consequence of patching the _autoincrement_column attribute. It doesn't seem to cause any issues, but ¯_(ツ)_/¯

user_53209060
Related