diff --git a/metrics/timescaledb/db.py b/metrics/timescaledb/db.py index ac284cf..02e2244 100644 --- a/metrics/timescaledb/db.py +++ b/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 diff --git a/tests/metrics/timescaledb/test_db.py b/tests/metrics/timescaledb/test_db.py index a16976e..ee2fca8 100644 --- a/tests/metrics/timescaledb/test_db.py +++ b/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)