Prepared psqlpy for OTLP

Signed-off-by: chandr-andr (Kiselev Aleksandr) <chandr@chandr.net>

Author: chandr-andr

python/tests/test_binary_copy.py

@@ -47,11 +47,10 @@ async def test_binary_copy_to_table_in_connection(
     buf.write(encoder.finish())
     buf.seek(0)
 
-    async with psql_pool.acquire() as connection:
-        inserted_rows = await connection.binary_copy_to_table(
-            source=buf,
-            table_name=table_name,
-        )
+    inserted_rows = await connection.binary_copy_to_table(
+        source=buf,
+        table_name=table_name,
+    )
 
     expected_inserted_row: typing.Final = 32
 

python/tests/test_listener.py

@@ -62,7 +62,9 @@ async def notify(
     if with_delay:
         await asyncio.sleep(0.5)
 
-    await (await psql_pool.connection()).execute(f"NOTIFY {channel}, '{TEST_PAYLOAD}'")
+    connection = await psql_pool.connection()
+    await connection.execute(f"NOTIFY {channel}, '{TEST_PAYLOAD}'")
+    connection.back_to_pool()
 
 
 async def check_insert_callback(
@@ -71,8 +73,9 @@ async def check_insert_callback(
     is_insert_exist: bool = True,
     number_of_data: int = 1,
 ) -> None:
+    connection = await psql_pool.connection()
     test_data_seq = (
-        await (await psql_pool.connection()).execute(
+        await connection.execute(
             f"SELECT * FROM {listener_table_name}",
         )
     ).result()
@@ -88,14 +91,18 @@ async def check_insert_callback(
     assert data_record["payload"] == TEST_PAYLOAD
     assert data_record["channel"] == TEST_CHANNEL
 
+    connection.back_to_pool()
+
 
 async def clear_test_table(
     psql_pool: ConnectionPool,
     listener_table_name: str,
 ) -> None:
-    await (await psql_pool.connection()).execute(
+    connection = await psql_pool.connection()
+    await connection.execute(
         f"DELETE FROM {listener_table_name}",
     )
+    connection.back_to_pool()
 
 
 @pytest.mark.usefixtures("create_table_for_listener_tests")
@@ -244,7 +251,8 @@ async def test_listener_more_than_one_callback(
         number_of_data=2,
     )
 
-    query_result = await (await psql_pool.connection()).execute(
+    connection = await psql_pool.connection()
+    query_result = await connection.execute(
         querystring=(f"SELECT * FROM {listener_table_name} WHERE channel = $1"),
         parameters=(additional_channel,),
     )

python/tests/test_transaction.py

@@ -39,8 +39,7 @@ async def test_transaction_init_parameters(
     deferrable: bool | None,
     read_variant: ReadVariant | None,
 ) -> None:
-    connection = await psql_pool.connection()
-    async with connection.transaction(
+    async with psql_pool.acquire() as connection, connection.transaction(
         isolation_level=isolation_level,
         deferrable=deferrable,
         read_variant=read_variant,
@@ -79,6 +78,8 @@ async def test_transaction_begin(
 
     assert len(result.result()) == number_database_records
 
+    await transaction.commit()
+
 
 async def test_transaction_commit(
     psql_pool: ConnectionPool,
@@ -97,15 +98,16 @@ async def test_transaction_commit(
 
     # Make request from other connection, it mustn't know
     # about new INSERT data before commit.
-    result = await (await psql_pool.connection()).execute(
+    connection = await psql_pool.connection()
+    result = await connection.execute(
         f"SELECT * FROM {table_name} WHERE name = $1",
         parameters=[test_name],
     )
     assert not result.result()
 
     await transaction.commit()
 
-    result = await (await psql_pool.connection()).execute(
+    result = await connection.execute(
         f"SELECT * FROM {table_name} WHERE name = $1",
         parameters=[test_name],
     )
@@ -136,7 +138,8 @@ async def test_transaction_savepoint(
     assert result.result()
 
     await transaction.rollback_savepoint(savepoint_name=savepoint_name)
-    result = await (await psql_pool.connection()).execute(
+    connection = await psql_pool.connection()
+    result = await connection.execute(
         f"SELECT * FROM {table_name} WHERE name = $1",
         parameters=[test_name],
     )
@@ -174,10 +177,12 @@ async def test_transaction_rollback(
             parameters=[test_name],
         )
 
-    result_from_conn = await (await psql_pool.connection()).execute(
+    connection = await psql_pool.connection()
+    result_from_conn = await connection.execute(
         f"INSERT INTO {table_name} VALUES ($1, $2)",
         parameters=[100, test_name],
     )
+    connection.back_to_pool()
 
     assert not (result_from_conn.result())
 
@@ -344,14 +349,17 @@ async def test_transaction_send_underlying_connection_to_pool_manually(
 
 async def test_execute_batch_method(psql_pool: ConnectionPool) -> None:
     """Test `execute_batch` method."""
-    await (await psql_pool.connection()).execute(querystring="DROP TABLE IF EXISTS execute_batch")
-    await (await psql_pool.connection()).execute(querystring="DROP TABLE IF EXISTS execute_batch2")
+    connection = await psql_pool.connection()
+    await connection.execute(querystring="DROP TABLE IF EXISTS execute_batch")
+    await connection.execute(querystring="DROP TABLE IF EXISTS execute_batch2")
     query = "CREATE TABLE execute_batch (name VARCHAR);CREATE TABLE execute_batch2 (name VARCHAR);"
-    async with psql_pool.acquire() as conn, conn.transaction() as transaction:
+    async with connection.transaction() as transaction:
         await transaction.execute_batch(querystring=query)
         await transaction.execute(querystring="SELECT * FROM execute_batch")
         await transaction.execute(querystring="SELECT * FROM execute_batch2")
 
+    connection.back_to_pool()
+
 
 @pytest.mark.parametrize(
     "synchronous_commit",