`Session.handle_query()` has no way to return `affected_rows` or `last_insert_id` for non-SELECT statements

[tarapt]

Context: I am using mysql-mimic as proxy to restrict certain queries while connecting to my organization's database on google cloud. I used claude to create the script and came across this bug. It's possible, that there is a simpler workaround, and claude couldn't find it. I don't think the below is a serious issue, but claude prepared this after some bug fixes, so I am reporting it for review incase it's important.

Summary

Session.handle_query returns one of ResultSet, (rows, columns), or None. There is no slot for affected_rows or last_insert_id. For non-SELECT statements, Connection.ok() always sends an OK packet with both fields = 0. This breaks IDE clients and ORMs that rely on update counts and auto-increment values.

Versions

• mysql-mimic==2.6.4
• Client: SQLAlchemy 1.4 + mysqlclient 2.2.x; JetBrains DataGrip / PyCharm Database Tools

Use case

A passthrough proxy that forwards SQL to a real MySQL backend and wants to forward cursor.rowcount and cursor.lastrowid from the backend back to the client unchanged.

Source

In mysql_mimic/connection.py:

async def _com_query(self, ...):
    result_set = await self.query(com_query.sql, com_query.query_attrs)
    if not result_set:
        await self.stream.write(self.ok())   # ← always 0/0
        return
    await self.write_text_resultset(result_set)

def ok(self, **kwargs):
    return packets.make_ok(
        capabilities=self.capabilities,
        status_flags=self.status_flags,
        **kwargs,                             # affected_rows/last_insert_id default to 0
    )

The implementer's Session has no way to influence what's in that OK packet for non-SELECT results.

Why it matters

JetBrains DataGrip / PyCharm Database Tools treats affected_rows=0 from an UPDATE as a row-vanished safety violation:

com.intellij.database.console.JdbcEngineUtils$UnsafeUpdateRolledBackException:
  Unexpected update count received (Actual: 0, Expected: 1).
  All changes will be rolled back.

In-grid edits that would have committed cleanly are refused.

ORMs that rely on cursor.rowcount (e.g. SQLAlchemy's update_match_count, optimistic-locking checks) silently break for the same reason. Anything reading cursor.lastrowid after INSERT gets 0.

Reproduction

A passthrough Session.handle_query that runs cursor.execute("UPDATE …") on a real backend connection where the row exists. SQLAlchemy:

import sqlalchemy as sa
e = sa.create_engine("mysql+mysqldb://user@proxy_host:proxy_port/db")
with e.connect() as c:
    r = c.execute(sa.text("UPDATE t SET val=99 WHERE id=1"))
    print(r.rowcount)        # → 0 even though one row was actually updated
    r = c.execute(sa.text("INSERT INTO t (val) VALUES (1)"))
    print(r.lastrowid)       # → 0 regardless of the auto-increment value

Suggested fix

Backwards-compatible: extend ResultSet with optional fields, and have Connection._com_query plumb them into the OK packet when present:

@dataclass
class ResultSet:
    rows: Iterable[Sequence] | AsyncIterable[Sequence] = ()
    columns: Sequence[ResultColumn] = ()
    affected_rows: int | None = None     # new
    last_insert_id: int | None = None    # new

Existing implementations that return (rows, columns) or None continue to work unchanged. Implementers who care about write metadata return a ResultSet with the new fields populated.

Workaround

Module-level monkey-patch of Connection.ok to read pending values stashed by the session via a _pending_ok_kwargs attribute. Works on 2.6.x but obviously fragile across versions:

from mysql_mimic.connection import Connection
_orig_ok = Connection.ok
def _ok_with_pending(self, **kwargs):
    pending = getattr(self.session, "_pending_ok_kwargs", None)
    if pending:
        kwargs = {**kwargs, **pending}
        self.session._pending_ok_kwargs = None
    return _orig_ok(self, **kwargs)
Connection.ok = _ok_with_pending

Then in handle_query, after a non-SELECT:

self._pending_ok_kwargs = {
    "affected_rows": cursor.rowcount or 0,
    "last_insert_id": cursor.lastrowid or 0,
}

[barakalon]

Mysql-mimic was built first for a read-proxy for an analytical systems, so there will likely be some rough edges like this when implementing a "passthrough proxy". But contributions are welcome :)

The suggested fix looks good to me.