From 841b39f7c5e6e11d754692f8772f025a4ad5d37a Mon Sep 17 00:00:00 2001 From: Alexey Masterov Date: Tue, 10 Sep 2024 12:52:46 +0200 Subject: [PATCH] Some refactoring --- .../cloud_regress/test_cloud_regress.py | 20 ++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/test_runner/cloud_regress/test_cloud_regress.py b/test_runner/cloud_regress/test_cloud_regress.py index 3b6b6a5c3c..5ab6fafad7 100644 --- a/test_runner/cloud_regress/test_cloud_regress.py +++ b/test_runner/cloud_regress/test_cloud_regress.py @@ -52,15 +52,17 @@ def setup(remote_pg: RemotePostgres): "RETURNS int AS 'regress.so' LANGUAGE C STRICT STABLE PARALLEL SAFE;" ) conn.rollback() - yield - log.info("Looking for extra roles...") - #cur.execute( - # "SELECT rolname FROM pg_catalog.pg_roles WHERE oid > 16384 AND rolname <> 'neondb_owner'" - #) - #log.info("Rows count: %s", cur.rowcount) - #for role in cur: - # cur.execute(f"DROP ROLE {role[0]}") - #conn.commit() + yield + log.info("Looking for extra roles...") + with psycopg2.connect(remote_pg.connstr()) as conn: + with conn.cursor() as cur: + cur.execute( + "SELECT rolname FROM pg_catalog.pg_roles WHERE oid > 16384 AND rolname <> 'neondb_owner'" + ) + log.info("Rows count: %s", cur.rowcount) + for role in cur: + cur.execute(f"DROP ROLE {role[0]}") + conn.commit() @pytest.mark.timeout(7200)