Move explanatory comment into function

metrics/timescaledb/db.py

@@ -29,15 +29,15 @@ def write(table, rows):
 
 
 def upsert(table, rows):
-    _ensure_table(table)
-    batch_size = _batch_size(table)
-    non_pk_columns = set(table.columns) - set(table.primary_key.columns)
+    with _get_engine().begin() as connection:
+        _ensure_table(connection, table)
+        batch_size = _batch_size(table)
+        non_pk_columns = set(table.columns) - set(table.primary_key.columns)
 
-    # use the primary key constraint to match rows to be updated,
-    # using default constraint name if not otherwise specified
-    constraint = table.primary_key.name or table.name + "_pkey"
+        # use the primary key constraint to match rows to be updated,
+        # using default constraint name if not otherwise specified
+        constraint = table.primary_key.name or table.name + "_pkey"
 
-    with _get_engine().begin() as connection:
         for values in batched(rows, batch_size):
             # Construct a PostgreSQL "INSERT..ON CONFLICT" style upsert statement
             # https://docs.sqlalchemy.org/en/20/dialects/postgresql.html#insert-on-conflict-upsert

tests/metrics/timescaledb/test_db.py

@@ -35,6 +35,9 @@ def get_rows(engine, table):
 
 
 def assert_is_hypertable(connection, engine, table):
+    # check there are timescaledb child tables
+    # https://stackoverflow.com/questions/1461722/how-to-find-child-tables-that-inherit-from-another-table-in-psql
+
     sql = """
     SELECT
       count(*)
@@ -83,9 +86,6 @@ def test_ensure_hypertable(engine, hypertable):
         assert not db._has_table(connection, hypertable)
         db._ensure_table(connection, hypertable)
         assert db._has_table(connection, hypertable)
-
-        # check there are timescaledb child tables
-        # https://stackoverflow.com/questions/1461722/how-to-find-child-tables-that-inherit-from-another-table-in-psql
         assert_is_hypertable(connection, engine, hypertable)