From 9b13728072944bba600ad3e21cafab42306698ba Mon Sep 17 00:00:00 2001
From: Brian Calvert <brian@graft.com>
Date: Tue, 10 Sep 2024 09:15:50 -0700
Subject: [PATCH] Use Engine.begin() instead of Engine.connect() to ensure
 DML/DDL statements are committed in migration code (#1129)

# Description
One of the changes in SQLAlchemy from 1.4.X -> 2.X is that they've
removed automatic commits for DDL/DML statements fed into `Connection`s
created from `Engine.connect()`. This leads to a few minor bugs in the
DB migration code post `0.39.0` release (and the SQLAlchemy 2.X
migration in #1127).

## This PR
This PR fixes those bugs by using `Connection`s created by
`Engine.begin()` instead

### Testing
```bash
$ bazel run sematic/db/tests:test_migrate
```
+ manual testing in a repo that's using Sematic as a dependency
---
 sematic/db/migrate.py                                         | 2 +-
 .../db/migrations/20220930014400_migrate_exception_data.py    | 2 +-
 sematic/db/tests/test_migrate.py                              | 4 ++--
 3 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/sematic/db/migrate.py b/sematic/db/migrate.py
index 40067ec1..a7aecdb5 100644
--- a/sematic/db/migrate.py
+++ b/sematic/db/migrate.py
@@ -51,7 +51,7 @@ def _is_migration_file(file_name: str) -> bool:
 
 
 def _get_current_versions() -> List[str]:
-    with db().get_engine().connect() as conn:
+    with db().get_engine().begin() as conn:
         conn.execute(
             text(
                 "CREATE TABLE IF NOT EXISTS "
diff --git a/sematic/db/migrations/20220930014400_migrate_exception_data.py b/sematic/db/migrations/20220930014400_migrate_exception_data.py
index fddb4b2c..8f93efaf 100644
--- a/sematic/db/migrations/20220930014400_migrate_exception_data.py
+++ b/sematic/db/migrations/20220930014400_migrate_exception_data.py
@@ -30,7 +30,7 @@ def up():
 
 
 def down():
-    with db().get_engine().connect() as conn:
+    with db().get_engine().begin() as conn:
         # TODO #303: standardize NULL vs 'null'
         run_id_exception_json_pairs = conn.execute(
             text(
diff --git a/sematic/db/tests/test_migrate.py b/sematic/db/tests/test_migrate.py
index 793bf998..b7ccf1c3 100644
--- a/sematic/db/tests/test_migrate.py
+++ b/sematic/db/tests/test_migrate.py
@@ -83,7 +83,7 @@ def test_invalid_sql(_, test_db_empty):  # noqa: F811
 def test_migrate(_, test_db_empty):  # noqa: F811
 
     with pytest.raises(OperationalError):
-        with db().get_engine().connect() as conn:
+        with db().get_engine().begin() as conn:
             conn.execute(text("SELECT version FROM schema_migrations;"))
 
     migrate_up()
@@ -93,7 +93,7 @@ def test_migrate(_, test_db_empty):  # noqa: F811
     assert len(current_versions) > 0
 
     # Test tables were created
-    with db().get_engine().connect() as conn:
+    with db().get_engine().begin() as conn:
         run_count = conn.execute(text("SELECT COUNT(*) from runs;"))
 
     assert list(run_count)[0][0] == 0