From aecf6584ea5489624ac9ed3df56045f5bf451a86 Mon Sep 17 00:00:00 2001 From: umd0730 <187863547+umd0730@users.noreply.github.com> Date: Mon, 21 Sep 2026 14:16:06 +0900 Subject: [PATCH] Commit writes made through owned SQL connections --- agatesql/table.py | 64 ++++++++++++---------- docs/index.rst | 20 +++++++ tests/test_transactions.py | 107 +++++++++++++++++++++++++++++++++++++ 3 files changed, 163 insertions(+), 28 deletions(-) create mode 100644 tests/test_transactions.py diff --git a/agatesql/table.py b/agatesql/table.py index 86daa6f..8d34831 100644 --- a/agatesql/table.py +++ b/agatesql/table.py @@ -5,6 +5,7 @@ import datetime import decimal +from contextlib import nullcontext from urllib.parse import urlsplit import agate @@ -248,7 +249,10 @@ def to_sql(self, connection_or_string, table_name, overwrite=False, Monkey patched as instance method :meth:`Table.to_sql`. :param connection_or_string: - An existing sqlalchemy connection or a connection string. + An existing sqlalchemy connection or a connection string. Writes using + a connection string are committed on success and rolled back on error. + An existing connection remains open and its transaction is managed by + the caller. :param table_name: The name of the SQL table to create. :param overwrite: @@ -276,34 +280,38 @@ def to_sql(self, connection_or_string, table_name, overwrite=False, """ engine, connection = get_engine_and_connection(connection_or_string) - dialect = connection.engine.dialect.name - sql_table = make_sql_table(self, table_name, dialect=dialect, db_schema=db_schema, constraints=constraints, - unique_constraint=unique_constraint, connection=connection, - min_col_len=min_col_len, col_len_multiplier=col_len_multiplier) - - if create: - if overwrite: - sql_table.drop(bind=connection, checkfirst=True) - - sql_table.create(bind=connection, checkfirst=create_if_not_exists) - - if insert: - insert = sql_table.insert() - for prefix in prefixes: - insert = insert.prefix_with(prefix) - if chunk_size is None: - connection.execute(insert, [dict(zip(self.column_names, row)) for row in self.rows]) - else: - number_of_rows = len(self.rows) - for index in range((number_of_rows - 1) // chunk_size + 1): - end_index = (index + 1) * chunk_size - if end_index > number_of_rows: - end_index = number_of_rows - connection.execute(insert, [dict(zip(self.column_names, row)) for row in - self.rows[index * chunk_size:end_index]]) - try: - return sql_table + # Only manage transactions for connections opened by this function. + # In particular, csvkit's csvsql manages its own outer transaction. + with connection.begin() if engine is not None else nullcontext(): + dialect = connection.engine.dialect.name + sql_table = make_sql_table( + self, table_name, dialect=dialect, db_schema=db_schema, constraints=constraints, + unique_constraint=unique_constraint, connection=connection, + min_col_len=min_col_len, col_len_multiplier=col_len_multiplier) + + if create: + if overwrite: + sql_table.drop(bind=connection, checkfirst=True) + + sql_table.create(bind=connection, checkfirst=create_if_not_exists) + + if insert: + insert = sql_table.insert() + for prefix in prefixes: + insert = insert.prefix_with(prefix) + if chunk_size is None: + connection.execute(insert, [dict(zip(self.column_names, row)) for row in self.rows]) + else: + number_of_rows = len(self.rows) + for index in range((number_of_rows - 1) // chunk_size + 1): + end_index = (index + 1) * chunk_size + if end_index > number_of_rows: + end_index = number_of_rows + connection.execute(insert, [dict(zip(self.column_names, row)) for row in + self.rows[index * chunk_size:end_index]]) + + return sql_table finally: if engine is not None: connection.close() diff --git a/docs/index.rst b/docs/index.rst index 04a369c..c288867 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -47,6 +47,26 @@ The first argument to either function can be any valid `sqlalchemy connection st That's all there is to it. +Transactions +------------ + +When passed a connection string, :meth:`.to_sql` commits its writes on success +and rolls them back on error, then closes its connection. When passed an existing +SQLAlchemy connection, it leaves transaction management and connection cleanup +to the caller. For example, to commit multiple writes together: + +.. code-block:: python + + from sqlalchemy import create_engine + + engine = create_engine('postgresql:///hospitals') + with engine.begin() as connection: + doctors.to_sql(connection, 'doctors') + departments.to_sql(connection, 'departments') + +Database-specific restrictions on transactional DDL still apply when creating +or replacing tables. + === API === diff --git a/tests/test_transactions.py b/tests/test_transactions.py new file mode 100644 index 0000000..375ffda --- /dev/null +++ b/tests/test_transactions.py @@ -0,0 +1,107 @@ +from unittest.mock import patch + +import agate +import pytest +from sqlalchemy import create_engine, text +from sqlalchemy.exc import IntegrityError + +import agatesql + + +@pytest.fixture +def database_url(tmp_path): + return 'sqlite:///' + (tmp_path / 'transactions.db').as_posix() + + +def make_table(rows): + return agate.Table(rows, ['id', 'name'], [agate.Number(), agate.Text()]) + + +def read_rows(url): + engine = create_engine(url) + try: + with engine.connect() as connection: + return [tuple(row) for row in connection.execute(text('SELECT id, name FROM people ORDER BY id'))] + finally: + engine.dispose() + + +@pytest.mark.parametrize('chunk_size', [None, 1]) +def test_to_sql_commits_owned_connection(database_url, chunk_size): + table = make_table([(1, 'Alice'), (2, 'Bob')]) + + table.to_sql(database_url, 'people', chunk_size=chunk_size) + + assert read_rows(database_url) == [(1, 'Alice'), (2, 'Bob')] + + +@pytest.mark.parametrize('chunk_size', [None, 1]) +def test_to_sql_rolls_back_failed_insert_and_closes_owned_connection(database_url, chunk_size): + engine = create_engine(database_url) + with engine.begin() as connection: + connection.execute(text('CREATE TABLE people (id INTEGER UNIQUE, name TEXT)')) + connection.execute(text("INSERT INTO people VALUES (1, 'original')")) + + connection = engine.connect() + table = make_table([(2, 'new'), (1, 'duplicate')]) + try: + with patch.object(agatesql.table, 'get_engine_and_connection', return_value=(engine, connection)): + with patch.object(engine, 'dispose', wraps=engine.dispose) as dispose: + with pytest.raises(IntegrityError): + table.to_sql(database_url, 'people', create=False, chunk_size=chunk_size) + assert connection.closed + dispose.assert_called_once_with() + + assert read_rows(database_url) == [(1, 'original')] + finally: + connection.close() + engine.dispose() + + +@pytest.mark.parametrize('commit', [False, True]) +def test_to_sql_preserves_caller_transaction(database_url, commit): + engine = create_engine(database_url) + try: + with engine.begin() as connection: + connection.execute(text('CREATE TABLE people (id INTEGER, name TEXT)')) + + with engine.connect() as connection: + transaction = connection.begin() + table = make_table([(1, 'Alice'), (2, 'Bob')]) + table.to_sql(connection, 'people', create=False, chunk_size=1) + + assert not connection.closed + assert transaction.is_active + assert connection.execute(text('SELECT COUNT(*) FROM people')).scalar() == 2 + assert read_rows(database_url) == [] + if commit: + transaction.commit() + else: + transaction.rollback() + + assert read_rows(database_url) == ([(1, 'Alice'), (2, 'Bob')] if commit else []) + finally: + engine.dispose() + + +def test_to_sql_preserves_caller_transaction_on_failure(database_url): + engine = create_engine(database_url) + try: + with engine.begin() as connection: + connection.execute(text('CREATE TABLE people (id INTEGER UNIQUE, name TEXT)')) + + with engine.connect() as connection: + transaction = connection.begin() + connection.execute(text("INSERT INTO people VALUES (1, 'caller')")) + table = make_table([(2, 'new'), (1, 'duplicate')]) + with pytest.raises(IntegrityError): + table.to_sql(connection, 'people', create=False, chunk_size=1) + + assert not connection.closed + assert transaction.is_active + assert connection.execute(text('SELECT COUNT(*) FROM people')).scalar() == 2 + transaction.rollback() + + assert read_rows(database_url) == [] + finally: + engine.dispose()