#!/usr/bin/env python
# cardinal_pythonlib/sqlalchemy/insert_on_duplicate.py
"""
===============================================================================
Original code copyright (C) 2009-2022 Rudolf Cardinal (rudolf@pobox.com).
This file is part of cardinal_pythonlib.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
===============================================================================
**Add "INSERT ON DUPLICATE KEY UPDATE" functionality to SQLAlchemy for MySQL.**
- https://www.reddit.com/r/Python/comments/p5grh/sqlalchemy_whats_the_idiomatic_way_of_writing/
- https://github.com/bedwards/sqlalchemy_mysql_ext/blob/master/duplicate.py
... modified
- https://docs.sqlalchemy.org/en/rel_1_0/core/compiler.html
- https://stackoverflow.com/questions/6611563/sqlalchemy-on-duplicate-key-update
- https://dev.mysql.com/doc/refman/5.7/en/insert-on-duplicate.html
Once implemented, you can do
.. code-block:: python
q = sqla_table.insert_on_duplicate().values(destvalues)
session.execute(q)
**Note: superseded by SQLAlchemy v1.2:**
- https://docs.sqlalchemy.org/en/latest/changelog/migration_12.html
- https://docs.sqlalchemy.org/en/latest/dialects/mysql.html#mysql-insert-on-duplicate-key-update
""" # noqa
import re
from typing import Any
from sqlalchemy.ext.compiler import compiles
from sqlalchemy.sql.compiler import SQLCompiler
from sqlalchemy.sql.expression import Insert, TableClause
from cardinal_pythonlib.logs import get_brace_style_log_with_null_handler
from cardinal_pythonlib.sqlalchemy.dialect import SqlaDialectName
log = get_brace_style_log_with_null_handler(__name__)
# noinspection PyAbstractClass
[docs]class InsertOnDuplicate(Insert):
"""
Class that derives from :class:`Insert`, so we can hook in to operations
involving it.
"""
pass
[docs]def insert_on_duplicate(
tablename: str, values: Any = None, inline: bool = False, **kwargs
):
"""
Command to produce an :class:`InsertOnDuplicate` object.
Args:
tablename: name of the table
values: values to ``INSERT``
inline: as per
https://docs.sqlalchemy.org/en/latest/core/dml.html#sqlalchemy.sql.expression.insert
kwargs: additional parameters
Returns:
an :class:`InsertOnDuplicate` object
""" # noqa
return InsertOnDuplicate(tablename, values, inline=inline, **kwargs)
# noinspection PyPep8Naming
[docs]def monkeypatch_TableClause() -> None:
"""
Modifies :class:`sqlalchemy.sql.expression.TableClause` to insert
a ``insert_on_duplicate`` member that is our :func:`insert_on_duplicate`
function as above.
"""
log.debug(
"Adding 'INSERT ON DUPLICATE KEY UPDATE' support for MySQL "
"to SQLAlchemy"
)
TableClause.insert_on_duplicate = insert_on_duplicate
# noinspection PyPep8Naming
[docs]def unmonkeypatch_TableClause() -> None:
"""
Reverses the action of :func:`monkeypatch_TableClause`.
"""
del TableClause.insert_on_duplicate
STARTSEPS = "`"
ENDSEPS = "`"
INSERT_FIELDNAMES_REGEX = (
r"^INSERT\sINTO\s[{startseps}]?(?P<table>\w+)[{endseps}]?\s+"
r"\((?P<columns>[, {startseps}{endseps}\w]+)\)\s+VALUES".format(
startseps=STARTSEPS, endseps=ENDSEPS
)
)
# http://pythex.org/ !
RE_INSERT_FIELDNAMES = re.compile(INSERT_FIELDNAMES_REGEX)
[docs]@compiles(InsertOnDuplicate, SqlaDialectName.MYSQL)
def compile_insert_on_duplicate_key_update(
insert: Insert, compiler: SQLCompiler, **kw
) -> str:
"""
Hooks into the use of the :class:`InsertOnDuplicate` class
for the MySQL dialect. Compiles the relevant SQL for an ``INSERT...
ON DUPLICATE KEY UPDATE`` statement.
Notes:
- We can't get the fieldnames directly from ``insert`` or ``compiler``.
- We could rewrite the innards of the visit_insert statement
(https://github.com/bedwards/sqlalchemy_mysql_ext/blob/master/duplicate.py)...
but, like that, it will get outdated.
- We could use a hack-in-by-hand method
(https://stackoverflow.com/questions/6611563/sqlalchemy-on-duplicate-key-update)
... but a little automation would be nice.
- So, regex to the rescue.
- NOTE THAT COLUMNS ARE ALREADY QUOTED by this stage; no need to repeat.
""" # noqa
# log.critical(compiler.__dict__)
# log.critical(compiler.dialect.__dict__)
# log.critical(insert.__dict__)
s = compiler.visit_insert(insert, **kw)
# log.critical(s)
m = RE_INSERT_FIELDNAMES.match(s)
if m is None:
raise ValueError("compile_insert_on_duplicate_key_update: no match")
columns = [c.strip() for c in m.group("columns").split(",")]
# log.critical(columns)
updates = ", ".join([f"{c} = VALUES({c})" for c in columns])
s += f" ON DUPLICATE KEY UPDATE {updates}"
# log.critical(s)
return s
_TEST_CODE = """
from sqlalchemy import Column, String, Integer, create_engine
from sqlalchemy.orm import Session
from sqlalchemy.ext.declarative import declarative_base
Base = declarative_base()
class OrmObject(Base):
__tablename__ = "sometable"
id = Column(Integer, primary_key=True)
name = Column(String)
engine = create_engine("sqlite://", echo=True)
Base.metadata.create_all(engine)
session = Session(engine)
d1 = dict(id=1, name="One")
d2 = dict(id=2, name="Two")
insert_1 = OrmObject.__table__.insert(values=d1)
insert_2 = OrmObject.__table__.insert(values=d2)
session.execute(insert_1)
session.execute(insert_2)
session.execute(insert_1) # raises sqlalchemy.exc.IntegrityError
# ... recommended cross-platform way is SELECT then INSERT or UPDATE
# accordingly; see
# https://groups.google.com/forum/#!topic/sqlalchemy/aQLqeHmLPQY
"""