test/alternator,cql-pytest: fix resource leak on failure

In the alternator and cql-pytest test frameworks, we have some convenient
contextmanager-based functions that allows us to create a temporary
resource (e.g., a table) that will be automatically deleted, for
example:

    with create_stream_test_table(...) as table:
        test_something(table)

However, our implementation of these functions wasn't safe. We had
code looking like:

    table = ...
    yield table
    table.delete()

The thinking was that the cleanup part (the table.delete()) will be
called after the user's code. However, if the user's code threw
(i.e., a failed assertion), the cleanup wasn't called... When the user's
code throws, it looks as if the "yield" throws. So the correct code
should look like:

    table = ...
    try:
        yield table
    finally:
        table.delete()

Python's contextmanager documentation indeed gives this idiom in its
example.

This patch fixes all contextmanager implementations in our tests to do
the cleanup even if the user's "with" block throws.

Signed-off-by: Nadav Har'El <nyh@scylladb.com>
Message-Id: <20210428083748.552203-1-nyh@scylladb.com>

test/alternator/test_streams.py

@@ -71,19 +71,22 @@ def create_stream_test_table(dynamodb, StreamViewType=None):
                     { 'AttributeName': 'p', 'AttributeType': 'S' },
                     { 'AttributeName': 'c', 'AttributeType': 'S' },
         ])
-    yield table
-    while True:
-        try:
-            table.delete()
-            return
-        except ClientError as ce:
-            # if the table has a stream currently being created we cannot
-            # delete the table immediately. Again, only with real dynamo
-            if ce.response['Error']['Code'] == 'ResourceInUseException':
-                print('Could not delete table yet. Sleeping 5s.')
-                time.sleep(5)
-                continue;
-            raise
+    try:
+        yield table
+    finally:
+        print(f"Deleting table {table.name}")
+        while True:
+            try:
+                table.delete()
+                break
+            except ClientError as ce:
+                # if the table has a stream currently being created we cannot
+                # delete the table immediately. Again, only with real dynamo
+                if ce.response['Error']['Code'] == 'ResourceInUseException':
+                    print('Could not delete table yet. Sleeping 5s.')
+                    time.sleep(5)
+                    continue;
+                raise
 
 def wait_for_active_stream(dynamodbstreams, table, timeout=60):
     exp = time.process_time() + timeout

test/cql-pytest/cassandra_tests/porting.py

@@ -46,23 +46,29 @@ def create_table(cql, keyspace, schema):
     table_name = previously_used_table_names.pop()
     table = keyspace + "." + table_name
     cql.execute("CREATE TABLE " + table + " " + schema)
-    yield table
-    cql.execute("DROP TABLE " + table)
-    previously_used_table_names.append(table_name)
+    try:
+        yield table
+    finally:
+        cql.execute("DROP TABLE " + table)
+        previously_used_table_names.append(table_name)
 
 @contextmanager
 def create_type(cql, keyspace, cmd):
     type_name = keyspace + "." + unique_name()
     cql.execute("CREATE TYPE " + type_name + " " + cmd)
-    yield type_name
-    cql.execute("DROP TYPE " + type_name)
+    try:
+        yield type_name
+    finally:
+        cql.execute("DROP TYPE " + type_name)
 
 @contextmanager
 def create_function(cql, keyspace, arg):
     function_name = keyspace + "." + unique_name()
     cql.execute("CREATE FUNCTION " + function_name + " " + arg)
-    yield function_name
-    cql.execute("DROP FUNCTION " + function_name)
+    try:
+        yield function_name
+    finally:
+        cql.execute("DROP FUNCTION " + function_name)
 
 # utility function to substitute table name in command.
 # For easy conversion of Java code, support %s as used in Java. Also support

test/cql-pytest/util.py

@@ -48,8 +48,10 @@ unique_name.last_ms = 0
 def new_test_keyspace(cql, opts):
     keyspace = unique_name()
     cql.execute("CREATE KEYSPACE " + keyspace + " " + opts)
-    yield keyspace
-    cql.execute("DROP KEYSPACE " + keyspace)
+    try:
+        yield keyspace
+    finally:
+        cql.execute("DROP KEYSPACE " + keyspace)
 
 # A utility function for creating a new temporary table with a given schema.
 # It can be used in a "with", as:
@@ -59,8 +61,10 @@ def new_test_keyspace(cql, opts):
 def new_test_table(cql, keyspace, schema, extra=""):
     table = keyspace + "." + unique_name()
     cql.execute("CREATE TABLE " + table + "(" + schema + ")" + extra)
-    yield table
-    cql.execute("DROP TABLE " + table)
+    try:
+        yield table
+    finally:
+        cql.execute("DROP TABLE " + table)
 
 def project(column_name_string, rows):
     """Returns a list of column values from each of the rows."""