From 3505f77a3a9ca31262eaa420425deea1aa2c5c24 Mon Sep 17 00:00:00 2001
From: Phillip Cloud <417981+cpcloud@users.noreply.github.com>
Date: Tue, 21 Jan 2025 07:33:05 -0500
Subject: [PATCH] test(fixtures): ensure that most connections are cleaned up
 when the session ends

---
 ibis/backends/clickhouse/tests/conftest.py |  9 +++++--
 ibis/backends/conftest.py                  |  4 +++-
 ibis/backends/datafusion/tests/conftest.py |  9 +++++--
 ibis/backends/duckdb/tests/conftest.py     |  9 +++++--
 ibis/backends/mssql/tests/conftest.py      |  9 +++++--
 ibis/backends/mysql/tests/conftest.py      |  9 +++++--
 ibis/backends/oracle/tests/conftest.py     |  9 +++++--
 ibis/backends/polars/tests/conftest.py     |  9 +++++--
 ibis/backends/postgres/tests/conftest.py   |  9 +++++--
 ibis/backends/pyspark/tests/conftest.py    | 28 +++++++++++++++-------
 ibis/backends/risingwave/tests/conftest.py |  9 +++++--
 ibis/backends/snowflake/tests/conftest.py  |  9 +++++--
 ibis/backends/sqlite/tests/conftest.py     |  9 +++++--
 ibis/backends/tests/base.py                |  6 +++++
 ibis/backends/tests/test_array.py          | 11 ++++-----
 ibis/backends/tests/test_client.py         |  4 ++--
 ibis/backends/tests/test_generic.py        |  4 ++--
 ibis/backends/tests/test_struct.py         |  8 +++----
 ibis/backends/trino/tests/conftest.py      |  9 +++++--
 19 files changed, 125 insertions(+), 48 deletions(-)

diff --git a/ibis/backends/clickhouse/tests/conftest.py b/ibis/backends/clickhouse/tests/conftest.py
index a2880a7e52e2..2edbea9680a5 100644
--- a/ibis/backends/clickhouse/tests/conftest.py
+++ b/ibis/backends/clickhouse/tests/conftest.py
@@ -178,8 +178,13 @@ def add_catalog_and_schema(node):
 
 
 @pytest.fixture(scope="session")
-def con(tmp_path_factory, data_dir, worker_id):
-    return TestConf.load_data(data_dir, tmp_path_factory, worker_id).connection
+def backend(data_dir, tmp_path_factory, worker_id):
+    return TestConf.load_data(data_dir, tmp_path_factory, worker_id)
+
+
+@pytest.fixture(scope="session")
+def con(backend):
+    return backend.connection
 
 
 @pytest.fixture(scope="session")
diff --git a/ibis/backends/conftest.py b/ibis/backends/conftest.py
index c1ca256bc638..f8111d90c0e0 100644
--- a/ibis/backends/conftest.py
+++ b/ibis/backends/conftest.py
@@ -470,7 +470,9 @@ def ddl_backend(request, data_dir, tmp_path_factory, worker_id):
 @pytest.fixture(scope="session")
 def ddl_con(ddl_backend):
     """Instance of Client, already connected to the db (if applies)."""
-    return ddl_backend.connection
+    connection = ddl_backend.connection
+    yield connection
+    connection.disconnect()
 
 
 @pytest.fixture(
diff --git a/ibis/backends/datafusion/tests/conftest.py b/ibis/backends/datafusion/tests/conftest.py
index 5dbdadecb907..8dc507aa50e3 100644
--- a/ibis/backends/datafusion/tests/conftest.py
+++ b/ibis/backends/datafusion/tests/conftest.py
@@ -74,8 +74,13 @@ def add_catalog_and_schema(node):
 
 
 @pytest.fixture(scope="session")
-def con(data_dir, tmp_path_factory, worker_id):
-    return TestConf.load_data(data_dir, tmp_path_factory, worker_id).connection
+def backend(data_dir, tmp_path_factory, worker_id):
+    return TestConf.load_data(data_dir, tmp_path_factory, worker_id)
+
+
+@pytest.fixture(scope="session")
+def con(backend):
+    return backend.connection
 
 
 @pytest.fixture(scope="session")
diff --git a/ibis/backends/duckdb/tests/conftest.py b/ibis/backends/duckdb/tests/conftest.py
index 629cfee2c9d1..2032120fbefe 100644
--- a/ibis/backends/duckdb/tests/conftest.py
+++ b/ibis/backends/duckdb/tests/conftest.py
@@ -138,8 +138,13 @@ def add_catalog_and_schema(node):
 
 
 @pytest.fixture(scope="session")
-def con(data_dir, tmp_path_factory, worker_id):
-    return TestConf.load_data(data_dir, tmp_path_factory, worker_id).connection
+def backend(data_dir, tmp_path_factory, worker_id):
+    return TestConf.load_data(data_dir, tmp_path_factory, worker_id)
+
+
+@pytest.fixture(scope="session")
+def con(backend):
+    return backend.connection
 
 
 @pytest.fixture(scope="session")
diff --git a/ibis/backends/mssql/tests/conftest.py b/ibis/backends/mssql/tests/conftest.py
index 4dd86ed760da..18984acf476b 100644
--- a/ibis/backends/mssql/tests/conftest.py
+++ b/ibis/backends/mssql/tests/conftest.py
@@ -60,5 +60,10 @@ def connect(*, tmpdir, worker_id, **kw):
 
 
 @pytest.fixture(scope="session")
-def con(data_dir, tmp_path_factory, worker_id):
-    return TestConf.load_data(data_dir, tmp_path_factory, worker_id).connection
+def backend(data_dir, tmp_path_factory, worker_id):
+    return TestConf.load_data(data_dir, tmp_path_factory, worker_id)
+
+
+@pytest.fixture(scope="session")
+def con(backend):
+    return backend.connection
diff --git a/ibis/backends/mysql/tests/conftest.py b/ibis/backends/mysql/tests/conftest.py
index b8e4f5734d34..f862a22292d7 100644
--- a/ibis/backends/mysql/tests/conftest.py
+++ b/ibis/backends/mysql/tests/conftest.py
@@ -75,5 +75,10 @@ def connect(*, tmpdir, worker_id, **kw):
 
 
 @pytest.fixture(scope="session")
-def con(tmp_path_factory, data_dir, worker_id):
-    return TestConf.load_data(data_dir, tmp_path_factory, worker_id).connection
+def backend(tmp_path_factory, data_dir, worker_id):
+    return TestConf.load_data(data_dir, tmp_path_factory, worker_id)
+
+
+@pytest.fixture(scope="session")
+def con(backend):
+    return backend.connection
diff --git a/ibis/backends/oracle/tests/conftest.py b/ibis/backends/oracle/tests/conftest.py
index adaa86c780f4..b8011c6f1bc3 100644
--- a/ibis/backends/oracle/tests/conftest.py
+++ b/ibis/backends/oracle/tests/conftest.py
@@ -133,8 +133,13 @@ def format_table(name: str) -> str:
 
 
 @pytest.fixture(scope="session")
-def con(data_dir, tmp_path_factory, worker_id):
-    return TestConf.load_data(data_dir, tmp_path_factory, worker_id).connection
+def backend(data_dir, tmp_path_factory, worker_id):
+    return TestConf.load_data(data_dir, tmp_path_factory, worker_id)
+
+
+@pytest.fixture(scope="session")
+def con(backend):
+    return backend.connection
 
 
 def init_oracle_database(
diff --git a/ibis/backends/polars/tests/conftest.py b/ibis/backends/polars/tests/conftest.py
index 8aec5127e305..9908b21a65af 100644
--- a/ibis/backends/polars/tests/conftest.py
+++ b/ibis/backends/polars/tests/conftest.py
@@ -44,8 +44,13 @@ def assert_series_equal(cls, left, right, *args, **kwargs) -> None:
 
 
 @pytest.fixture(scope="session")
-def con(data_dir, tmp_path_factory, worker_id):
-    return TestConf.load_data(data_dir, tmp_path_factory, worker_id).connection
+def backend(data_dir, tmp_path_factory, worker_id):
+    return TestConf.load_data(data_dir, tmp_path_factory, worker_id)
+
+
+@pytest.fixture(scope="session")
+def con(backend):
+    return backend.connection
 
 
 @pytest.fixture(scope="session")
diff --git a/ibis/backends/postgres/tests/conftest.py b/ibis/backends/postgres/tests/conftest.py
index a0ec63dfae20..fea7aaeebaea 100644
--- a/ibis/backends/postgres/tests/conftest.py
+++ b/ibis/backends/postgres/tests/conftest.py
@@ -69,8 +69,13 @@ def connect(*, tmpdir, worker_id, **kw):
 
 
 @pytest.fixture(scope="session")
-def con(tmp_path_factory, data_dir, worker_id):
-    return TestConf.load_data(data_dir, tmp_path_factory, worker_id).connection
+def backend(data_dir, tmp_path_factory, worker_id):
+    return TestConf.load_data(data_dir, tmp_path_factory, worker_id)
+
+
+@pytest.fixture(scope="session")
+def con(backend):
+    return backend.connection
 
 
 @pytest.fixture(scope="module")
diff --git a/ibis/backends/pyspark/tests/conftest.py b/ibis/backends/pyspark/tests/conftest.py
index 65df70f78142..35b402eaf3eb 100644
--- a/ibis/backends/pyspark/tests/conftest.py
+++ b/ibis/backends/pyspark/tests/conftest.py
@@ -11,6 +11,7 @@
 import pandas as pd
 import pytest
 from filelock import FileLock
+from pyspark.errors.exceptions.connect import SparkConnectException
 
 import ibis
 from ibis import util
@@ -378,11 +379,12 @@ def connect(*, tmpdir, worker_id, **kw):
             return con
 
     @pytest.fixture(scope="session")
-    def con_streaming(data_dir, tmp_path_factory, worker_id):
-        backend_test = TestConfForStreaming.load_data(
-            data_dir, tmp_path_factory, worker_id
-        )
-        return backend_test.connection
+    def backend_streaming(data_dir, tmp_path_factory, worker_id):
+        return TestConfForStreaming.load_data(data_dir, tmp_path_factory, worker_id)
+
+    @pytest.fixture(scope="session")
+    def con_streaming(backend_streaming):
+        return backend_streaming.connection
 
     @pytest.fixture(autouse=True, scope="function")
     def stop_active_jobs(con_streaming):
@@ -396,13 +398,21 @@ def write_to_memory(self, expr, table_name):
         df = self._session.sql(expr.compile())
         df.writeStream.format("memory").queryName(table_name).start()
 
+    def __del__(self):
+        try:  # noqa: SIM105
+            self.connection.disconnect()
+        except (AttributeError, SparkConnectException):
+            pass
+
 
 @pytest.fixture(scope="session")
-def con(data_dir, tmp_path_factory, worker_id):
-    backend_test = TestConf.load_data(data_dir, tmp_path_factory, worker_id)
-    con = backend_test.connection
+def backend(data_dir, tmp_path_factory, worker_id):
+    return TestConf.load_data(data_dir, tmp_path_factory, worker_id)
+
 
-    return con
+@pytest.fixture(scope="session")
+def con(backend):
+    return backend.connection
 
 
 class IbisWindow:
diff --git a/ibis/backends/risingwave/tests/conftest.py b/ibis/backends/risingwave/tests/conftest.py
index 06940a6573b6..bbcb94ffc66a 100644
--- a/ibis/backends/risingwave/tests/conftest.py
+++ b/ibis/backends/risingwave/tests/conftest.py
@@ -68,8 +68,13 @@ def connect(*, tmpdir, worker_id, port: int | None = None, **kw):
 
 
 @pytest.fixture(scope="session")
-def con(tmp_path_factory, data_dir, worker_id):
-    return TestConf.load_data(data_dir, tmp_path_factory, worker_id).connection
+def backend(tmp_path_factory, data_dir, worker_id):
+    return TestConf.load_data(data_dir, tmp_path_factory, worker_id)
+
+
+@pytest.fixture(scope="session")
+def con(backend):
+    return backend.connection
 
 
 @pytest.fixture(scope="module")
diff --git a/ibis/backends/snowflake/tests/conftest.py b/ibis/backends/snowflake/tests/conftest.py
index 8b309d71d7db..e80f75ec9dcd 100644
--- a/ibis/backends/snowflake/tests/conftest.py
+++ b/ibis/backends/snowflake/tests/conftest.py
@@ -215,5 +215,10 @@ def connect(*, tmpdir, worker_id, **kw) -> BaseBackend:
 
 
 @pytest.fixture(scope="session")
-def con(data_dir, tmp_path_factory, worker_id):
-    return TestConf.load_data(data_dir, tmp_path_factory, worker_id).connection
+def backend(data_dir, tmp_path_factory, worker_id):
+    return TestConf.load_data(data_dir, tmp_path_factory, worker_id)
+
+
+@pytest.fixture(scope="session")
+def con(backend):
+    return backend.connection
diff --git a/ibis/backends/sqlite/tests/conftest.py b/ibis/backends/sqlite/tests/conftest.py
index 7268fb802042..ecafdd4f6051 100644
--- a/ibis/backends/sqlite/tests/conftest.py
+++ b/ibis/backends/sqlite/tests/conftest.py
@@ -55,5 +55,10 @@ def functional_alltypes(self) -> ir.Table:
 
 
 @pytest.fixture(scope="session")
-def con(data_dir, tmp_path_factory, worker_id):
-    return TestConf.load_data(data_dir, tmp_path_factory, worker_id).connection
+def backend(data_dir, tmp_path_factory, worker_id):
+    return TestConf.load_data(data_dir, tmp_path_factory, worker_id)
+
+
+@pytest.fixture(scope="session")
+def con(backend):
+    return backend.connection
diff --git a/ibis/backends/tests/base.py b/ibis/backends/tests/base.py
index eed34fe4a473..4c4773259b41 100644
--- a/ibis/backends/tests/base.py
+++ b/ibis/backends/tests/base.py
@@ -63,6 +63,12 @@ class BackendTest(abc.ABC):
     tpc_absolute_tolerance: float | None = None
     "Absolute tolerance for floating point comparisons with pytest.approx in TPC correctness tests."
 
+    def __del__(self):
+        try:  # noqa: SIM105
+            self.connection.disconnect()
+        except AttributeError:
+            pass
+
     @property
     @abc.abstractmethod
     def deps(self) -> Iterable[str]:
diff --git a/ibis/backends/tests/test_array.py b/ibis/backends/tests/test_array.py
index 1bf9ce24a803..02f811bc3f0a 100644
--- a/ibis/backends/tests/test_array.py
+++ b/ibis/backends/tests/test_array.py
@@ -22,11 +22,10 @@
     GoogleBadRequest,
     MySQLOperationalError,
     PolarsComputeError,
-    PsycoPg2ArraySubscriptError,
     PsycoPg2IndeterminateDatatype,
     PsycoPg2InternalError,
     PsycoPg2ProgrammingError,
-    PsycoPg2SyntaxError,
+    PsycoPgInvalidTextRepresentation,
     PsycoPgSyntaxError,
     Py4JJavaError,
     PyAthenaDatabaseError,
@@ -1118,7 +1117,7 @@ def test_unnest_struct(con):
 
 
 @builtin_array
-@pytest.mark.notimpl(["postgres"], raises=PsycoPg2SyntaxError)
+@pytest.mark.notimpl(["postgres"], raises=PsycoPgSyntaxError)
 @pytest.mark.notimpl(["risingwave"], raises=PsycoPg2InternalError)
 @pytest.mark.notimpl(
     ["trino"], reason="inserting maps into structs doesn't work", raises=TrinoUserError
@@ -1209,7 +1208,7 @@ def test_zip_null(con, fn):
 
 
 @builtin_array
-@pytest.mark.notimpl(["postgres"], raises=PsycoPg2SyntaxError)
+@pytest.mark.notimpl(["postgres"], raises=PsycoPgSyntaxError)
 @pytest.mark.notimpl(["risingwave"], raises=PsycoPg2ProgrammingError)
 @pytest.mark.notimpl(["datafusion"], raises=Exception, reason="not yet supported")
 @pytest.mark.notimpl(
@@ -1769,7 +1768,7 @@ def test_table_unnest_column_expr(backend):
 @pytest.mark.notimpl(["datafusion", "polars"], raises=com.OperationNotDefinedError)
 @pytest.mark.notimpl(["trino"], raises=TrinoUserError)
 @pytest.mark.notimpl(["athena"], raises=PyAthenaOperationalError)
-@pytest.mark.notimpl(["postgres"], raises=PsycoPg2SyntaxError)
+@pytest.mark.notimpl(["postgres"], raises=PsycoPgSyntaxError)
 @pytest.mark.notimpl(["risingwave"], raises=PsycoPg2ProgrammingError)
 @pytest.mark.notyet(
     ["risingwave"], raises=PsycoPg2InternalError, reason="not supported in risingwave"
@@ -1890,7 +1889,7 @@ def test_array_agg_bool(con, data, agg, baseline_func):
 
 @pytest.mark.notyet(
     ["postgres"],
-    raises=PsycoPg2ArraySubscriptError,
+    raises=PsycoPgInvalidTextRepresentation,
     reason="all dimensions must match in size",
 )
 @pytest.mark.notimpl(["risingwave", "flink"], raises=com.OperationNotDefinedError)
diff --git a/ibis/backends/tests/test_client.py b/ibis/backends/tests/test_client.py
index bf645d5409f5..4664f927d133 100644
--- a/ibis/backends/tests/test_client.py
+++ b/ibis/backends/tests/test_client.py
@@ -32,7 +32,7 @@
     ImpalaHiveServer2Error,
     OracleDatabaseError,
     PsycoPg2InternalError,
-    PsycoPg2UndefinedObject,
+    PsycoPgUndefinedObject,
     Py4JJavaError,
     PyAthenaDatabaseError,
     PyODBCProgrammingError,
@@ -725,7 +725,7 @@ def test_list_database_contents(con):
 @pytest.mark.notyet(["databricks"], raises=DatabricksServerOperationError)
 @pytest.mark.notyet(["bigquery"], raises=com.UnsupportedBackendType)
 @pytest.mark.notyet(
-    ["postgres"], raises=PsycoPg2UndefinedObject, reason="no unsigned int types"
+    ["postgres"], raises=PsycoPgUndefinedObject, reason="no unsigned int types"
 )
 @pytest.mark.notyet(
     ["oracle"], raises=OracleDatabaseError, reason="no unsigned int types"
diff --git a/ibis/backends/tests/test_generic.py b/ibis/backends/tests/test_generic.py
index 98f48383ab42..7260656c0488 100644
--- a/ibis/backends/tests/test_generic.py
+++ b/ibis/backends/tests/test_generic.py
@@ -25,7 +25,7 @@
     OracleDatabaseError,
     PolarsInvalidOperationError,
     PsycoPg2InternalError,
-    PsycoPg2SyntaxError,
+    PsycoPgSyntaxError,
     Py4JJavaError,
     PyAthenaDatabaseError,
     PyAthenaOperationalError,
@@ -1735,7 +1735,7 @@ def hash_256(col):
                 pytest.mark.notimpl(["flink"], raises=Py4JJavaError),
                 pytest.mark.notimpl(["druid"], raises=PyDruidProgrammingError),
                 pytest.mark.notimpl(["oracle"], raises=OracleDatabaseError),
-                pytest.mark.notimpl(["postgres"], raises=PsycoPg2SyntaxError),
+                pytest.mark.notimpl(["postgres"], raises=PsycoPgSyntaxError),
                 pytest.mark.notimpl(["risingwave"], raises=PsycoPg2InternalError),
                 pytest.mark.notimpl(["snowflake"], raises=AssertionError),
                 pytest.mark.never(
diff --git a/ibis/backends/tests/test_struct.py b/ibis/backends/tests/test_struct.py
index 5801b97e52ce..48e046cb3a54 100644
--- a/ibis/backends/tests/test_struct.py
+++ b/ibis/backends/tests/test_struct.py
@@ -13,7 +13,7 @@
     DatabricksServerOperationError,
     PolarsColumnNotFoundError,
     PsycoPg2InternalError,
-    PsycoPg2SyntaxError,
+    PsycoPgSyntaxError,
     Py4JJavaError,
     PyAthenaDatabaseError,
     PyAthenaOperationalError,
@@ -138,7 +138,7 @@ def test_collect_into_struct(alltypes):
 
 
 @pytest.mark.notimpl(
-    ["postgres"], reason="struct literals not implemented", raises=PsycoPg2SyntaxError
+    ["postgres"], reason="struct literals not implemented", raises=PsycoPgSyntaxError
 )
 @pytest.mark.notimpl(
     ["risingwave"],
@@ -155,7 +155,7 @@ def test_field_access_after_case(con):
 
 
 @pytest.mark.notimpl(
-    ["postgres"], reason="struct literals not implemented", raises=PsycoPg2SyntaxError
+    ["postgres"], reason="struct literals not implemented", raises=PsycoPgSyntaxError
 )
 @pytest.mark.notimpl(["flink"], raises=IbisError, reason="not implemented in ibis")
 @pytest.mark.parametrize(
@@ -242,7 +242,7 @@ def test_keyword_fields(con, nullable):
 
 @pytest.mark.notyet(
     ["postgres"],
-    raises=PsycoPg2SyntaxError,
+    raises=PsycoPgSyntaxError,
     reason="sqlglot doesn't implement structs for postgres correctly",
 )
 @pytest.mark.notyet(
diff --git a/ibis/backends/trino/tests/conftest.py b/ibis/backends/trino/tests/conftest.py
index 3d9610d5f2d7..5987bc5ebb14 100644
--- a/ibis/backends/trino/tests/conftest.py
+++ b/ibis/backends/trino/tests/conftest.py
@@ -162,8 +162,13 @@ def awards_players(self):
 
 
 @pytest.fixture(scope="session")
-def con(tmp_path_factory, data_dir, worker_id):
-    return TestConf.load_data(data_dir, tmp_path_factory, worker_id).connection
+def backend(tmp_path_factory, data_dir, worker_id):
+    return TestConf.load_data(data_dir, tmp_path_factory, worker_id)
+
+
+@pytest.fixture(scope="session")
+def con(backend):
+    return backend.connection
 
 
 def generate_tpc_tables(suite_name, *, data_dir):