Commit writes made through owned SQL connections

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()

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
 ===

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()