Repository: ibis-project/ibis

Issue:
bug: Impala returns a cursor for operations that are expected to not return a cursor
### What happened?

Impala returns a cursor for operations that are expected to not return a cursor. This is a problem because ibis.backends.base.sql.raw_sql will release empty cursors and a lot of code relies on this behaviour, i.e. a lot of code doesn't capture the returned cursor and release it. For example, When you create a table this code (https://github.com/ibis-project/ibis/blob/master/ibis/backends/impala/__init__.py#L626) calls raw_sql and ignored the returned cursor. This violates the warning here: https://github.com/ibis-project/ibis/blob/master/ibis/backends/base/sql/__init__.py#L127

The cursor then goes out of scope and when garbage collection runs it will be closed and removed from the connection-pool - but in the meantime the connection-pool has to create new connections or can run out of connections if a lot of operations (that don't usually return cursors) are performed (e.g. during test setup/teardown when we are creating/dropping tables/databases). I.e. https://github.com/ibis-project/ibis/blob/2.1.1/ibis/backends/impala/client.py#L122 ( this is replaced in main by sqlalchemy QueuePool)

A hacky work-around is something like this - but it far from satisfactory:

```
class ImpalaCursor(ibis_impala_client.ImpalaCursor):
    def __bool__(self):
        # In ibis.backends.base.sql.raw_sql the code automatically releases any cursor that evaluates to False.
        # The (implied) intention is that operations that return no rows, like CREATE and DROP, should return
        # an empty cursor. Impala does return a cursor for these types of operations with one "row" that says
        # some like "The table has been created" (or similar). We can't examine the value of the rows fetched
        # as this would affect callers who wish to fetch the result - so instead we look for queries that
        # just return summary information.
        # REFRESH does indeed seem to return no cursor, hence the check for None
        # Note - will break for a query like
        #  SELECT CAST(COUNT(*) AS STRING) AS summary FROM __ibis_tmp.ibis_tmp_pandas_8bcdd4082ff0477ba4a11fb0a0cd12f4
        if (
            self._cursor is None
            or self._cursor.description is None
            or len(self._cursor.description) == 1
            and self._cursor.description[0]
            == (
                "summary",
                "STRING",
                None,
                None,
                None,
                None,
                None,
            )
        ):
            return False
        return True
```

It would be nice if there was some way to turn off these messages from Impala, but I couldn't find a way to do this. It may be related to the fact that sometimes these operations are long running and thus can return an asynchronous completion handle? (or whatever it is called)

### What version of ibis are you using?

2.1.1

### What backend(s) are you using, if any?

Impala

### Relevant log output

```sh
N/A - I can provide logging but and a traceback but it will be for 2.1.1 code, so the failure mode on
master will be different. The warning on https://github.com/ibis-project/ibis/blob/master/ibis/backends/base/sql/__init__.py#L127
is still being violated though by code like https://github.com/ibis-project/ibis/blob/2.1.1/ibis/backends/impala/client.py#L122
```


### Code of Conduct

- [X] I agree to follow this project's Code of Conduct


Discussion:
@Apteryx0 Thanks for the detailed report!

Anecdotally we've been seeing an uptick in pool exhaustion during impala test runs recently. Could be many reasons for that, but certainly properly cleaning up cursors would go some ways towards not exhausting the pool.

I'll put up a PR, but it may be difficult for you to back port to 2.x.

Resolving patch:
diff --git a/ibis/backends/base/sql/__init__.py b/ibis/backends/base/sql/__init__.py
--- a/ibis/backends/base/sql/__init__.py
+++ b/ibis/backends/base/sql/__init__.py
@@ -131,11 +131,7 @@ def raw_sql(self, query: str):
         query
             DDL or DML statement
         """
-        cursor = self.con.execute(query)
-        if cursor:
-            return cursor
-        cursor.release()
-        return None
+        return self.con.execute(query)
 
     @contextlib.contextmanager
     def _safe_raw_sql(self, *args, **kwargs):
diff --git a/ibis/backends/impala/__init__.py b/ibis/backends/impala/__init__.py
--- a/ibis/backends/impala/__init__.py
+++ b/ibis/backends/impala/__init__.py
@@ -8,6 +8,7 @@
 import os
 import re
 import weakref
+from functools import cached_property
 from pathlib import Path
 from posixpath import join as pjoin
 from typing import TYPE_CHECKING, Any, Literal
@@ -299,17 +300,15 @@ def do_connect(
 
         self._ensure_temp_db_exists()
 
-    @property
+    @cached_property
     def version(self):
-        cursor = self.raw_sql('select version()')
-        result = cursor.fetchone()[0]
-        cursor.release()
+        with self._safe_raw_sql('select version()') as cursor:
+            (result,) = cursor.fetchone()
         return result
 
     def list_databases(self, like=None):
-        cur = self.raw_sql('SHOW DATABASES')
-        databases = self._get_list(cur)
-        cur.release()
+        with self._safe_raw_sql('SHOW DATABASES') as cur:
+            databases = self._get_list(cur)
         return self._filter_with_like(databases, like)
 
     def list_tables(self, like=None, database=None):
@@ -317,25 +316,35 @@ def list_tables(self, like=None, database=None):
         if database is not None:
             statement += f' IN {database}'
         if like:
-            m = fully_qualified_re.match(like)
-            if m:
-                database, quoted, unquoted = m.groups()
+            if match := fully_qualified_re.match(like):
+                database, quoted, unquoted = match.groups()
                 like = quoted or unquoted
                 return self.list_tables(like=like, database=database)
             statement += f" LIKE '{like}'"
 
-        return self._filter_with_like(
-            [row[0] for row in self.raw_sql(statement).fetchall()]
-        )
+        with self._safe_raw_sql(statement) as cursor:
+            tables = [row[0] for row in cursor.fetchall()]
+        return self._filter_with_like(tables)
 
     def fetch_from_cursor(self, cursor, schema):
         batches = cursor.fetchall(columnar=True)
-        names = [x[0] for x in cursor.description]
+        names = [name for name, *_ in cursor.description]
         df = _column_batches_to_dataframe(names, batches)
         if schema:
             return schema.apply_to(df)
         return df
 
+    @contextlib.contextmanager
+    def _safe_raw_sql(self, *args, **kwargs):
+        with contextlib.closing(self.raw_sql(*args, **kwargs)) as cursor:
+            yield cursor
+        with contextlib.suppress(AttributeError):
+            cursor.release()
+
+    def _safe_exec_sql(self, *args, **kwargs):
+        with self._safe_raw_sql(*args, **kwargs):
+            pass
+
     @property
     def hdfs(self):
         if self._hdfs is None:
@@ -410,7 +419,7 @@ def create_database(self, name, path=None, force=False):
             # which is easier for manual cleanup, if necessary
             self.hdfs.mkdir(path)
         statement = CreateDatabase(name, path=path, can_exist=force)
-        return self.raw_sql(statement)
+        self._safe_exec_sql(statement)
 
     def drop_database(self, name, force=False):
         """Drop an Impala database.
@@ -433,7 +442,7 @@ def drop_database(self, name, force=False):
             udas = []
         if force:
             for table in tables:
-                util.log('Dropping {}'.format(f'{name}.{table}'))
+                util.log(f'Dropping {name}.{table}')
                 self.drop_table_or_view(table, database=name)
             for func in udfs:
                 util.log(f'Dropping function {func.name}({func.inputs})')
@@ -457,13 +466,9 @@ def drop_database(self, name, force=False):
                 "being dropped, or set force=True"
             )
         statement = DropDatabase(name, must_exist=not force)
-        return self.raw_sql(statement)
+        self._safe_exec_sql(statement)
 
-    def get_schema(
-        self,
-        table_name: str,
-        database: str | None = None,
-    ) -> sch.Schema:
+    def get_schema(self, table_name: str, database: str | None = None) -> sch.Schema:
         """Return a Schema object for the indicated table and database.
 
         Parameters
@@ -479,7 +484,7 @@ def get_schema(
             Ibis schema
         """
         qualified_name = self._fully_qualified_name(table_name, database)
-        query = f'DESCRIBE {qualified_name}'
+        query = f"DESCRIBE {qualified_name}"
 
         # only pull out the first two columns which are names and types
         pairs = [row[:2] for row in self.con.fetchall(query)]
@@ -493,9 +498,9 @@ def get_schema(
     def client_options(self):
         return self.con.options
 
-    def get_options(self):
+    def get_options(self) -> dict[str, str]:
         """Return current query options for the Impala session."""
-        return dict(row[:2] for row in self.con.fetchall("SET"))
+        return {key: value for key, value, *_ in self.con.fetchall("SET")}
 
     def set_options(self, options):
         self.con.set_options(options)
@@ -526,12 +531,12 @@ def create_view(
         ast = self.compiler.to_ast(obj)
         select = ast.queries[0]
         statement = CreateView(name, select, database=database, can_exist=overwrite)
-        self.raw_sql(statement)
+        self._safe_exec_sql(statement)
         return self.table(name, database=database)
 
     def drop_view(self, name, database=None, force=False):
-        statement = DropView(name, database=database, must_exist=not force)
-        return self.raw_sql(statement)
+        stmt = DropView(name, database=database, must_exist=not force)
+        self._safe_exec_sql(stmt)
 
     @contextlib.contextmanager
     def _setup_insert(self, obj):
@@ -609,7 +614,7 @@ def create_table(
 
                 if overwrite:
                     self.drop_table(name, force=True)
-                self.raw_sql(
+                self._safe_exec_sql(
                     CTAS(
                         name,
                         select,
@@ -623,7 +628,7 @@ def create_table(
         else:  # schema is not None
             if overwrite:
                 self.drop_table(name, force=True)
-            self.raw_sql(
+            self._safe_exec_sql(
                 CreateTableWithSchema(
                     name,
                     schema,
@@ -672,7 +677,7 @@ def avro_file(
         stmt = ddl.CreateTableAvro(
             name, hdfs_dir, avro_schema, database=database, external=external
         )
-        self.raw_sql(stmt)
+        self._safe_exec_sql(stmt)
         return self._wrap_new_table(name, database, persist)
 
     def delimited_file(
@@ -738,7 +743,7 @@ def delimited_file(
             lineterminator=lineterminator,
             escapechar=escapechar,
         )
-        self.raw_sql(stmt)
+        self._safe_exec_sql(stmt)
         return self._wrap_new_table(name, database, persist)
 
     def parquet_file(
@@ -821,7 +826,7 @@ def parquet_file(
             external=external,
             can_exist=False,
         )
-        self.raw_sql(stmt)
+        self._safe_exec_sql(stmt)
         return self._wrap_new_table(name, database, persist)
 
     def _get_concrete_table_path(self, name, database, persist=False):
@@ -870,7 +875,7 @@ def _wrap_new_table(self, name, database, persist):
                 qualified_name, cardinality
             )
         )
-        self.raw_sql(set_card)
+        self._safe_exec_sql(set_card)
 
         return t
 
@@ -931,7 +936,7 @@ def drop_table(
         >>> con.drop_table(table, database=db, force=True)  # doctest: +SKIP
         """
         statement = DropTable(name, database=database, must_exist=not force)
-        self.raw_sql(statement)
+        self._safe_exec_sql(statement)
 
     def truncate_table(self, name: str, database: str | None = None) -> None:
         """Delete all rows from an existing table.
@@ -944,7 +949,7 @@ def truncate_table(self, name: str, database: str | None = None) -> None:
             Database name
         """
         statement = TruncateTable(name, database=database)
-        self.raw_sql(statement)
+        self._safe_exec_sql(statement)
 
     def drop_table_or_view(self, name, *, database=None, force=False):
         """Drop view or table."""
@@ -976,14 +981,13 @@ def cache_table(self, table_name, *, database=None, pool='default'):
         >>> con.cache_table('my_table', database=db, pool=pool)  # doctest: +SKIP
         """
         statement = ddl.CacheTable(table_name, database=database, pool=pool)
-        self.raw_sql(statement)
+        self._safe_exec_sql(statement)
 
     def _get_schema_using_query(self, query):
-        cur = self.raw_sql(f"SELECT * FROM ({query}) t0 LIMIT 0")
-        # resets the state of the cursor and closes operation
-        cur.fetchall()
-        ibis_fields = self._adapt_types(cur.description)
-        cur.release()
+        with self._safe_raw_sql(f"SELECT * FROM ({query}) t0 LIMIT 0") as cur:
+            # resets the state of the cursor and closes operation
+            cur.fetchall()
+            ibis_fields = self._adapt_types(cur.description)
 
         return sch.Schema(ibis_fields)
 
@@ -1009,7 +1013,7 @@ def create_function(self, func, name=None, database=None):
             stmt = ddl.CreateUDA(func, name=name, database=database)
         else:
             raise TypeError(func)
-        self.raw_sql(stmt)
+        self._safe_exec_sql(stmt)
 
     def drop_udf(
         self,
@@ -1082,7 +1086,7 @@ def _drop_single_function(self, name, input_types, database=None, aggregate=Fals
             aggregate=aggregate,
             database=database,
         )
-        self.raw_sql(stmt)
+        self._safe_exec_sql(stmt)
 
     def _drop_all_functions(self, database):
         udfs = self.list_udfs(database=database)
@@ -1094,7 +1098,7 @@ def _drop_all_functions(self, database):
                 aggregate=False,
                 database=database,
             )
-            self.raw_sql(stmt)
+            self._safe_exec_sql(stmt)
         udafs = self.list_udas(database=database)
         for udaf in udafs:
             stmt = ddl.DropFunction(
@@ -1104,28 +1108,23 @@ def _drop_all_functions(self, database):
                 aggregate=True,
                 database=database,
             )
-            self.raw_sql(stmt)
+            self._safe_exec_sql(stmt)
 
     def list_udfs(self, database=None, like=None):
         """Lists all UDFs associated with given database."""
         if not database:
             database = self.current_database
         statement = ddl.ListFunction(database, like=like, aggregate=False)
-        cur = self.raw_sql(statement)
-        result = self._get_udfs(cur, udf.ImpalaUDF)
-        cur.release()
-        return result
+        with self._safe_raw_sql(statement) as cur:
+            return self._get_udfs(cur, udf.ImpalaUDF)
 
     def list_udas(self, database=None, like=None):
         """Lists all UDAFs associated with a given database."""
         if not database:
             database = self.current_database
         statement = ddl.ListFunction(database, like=like, aggregate=True)
-        cur = self.raw_sql(statement)
-        result = self._get_udfs(cur, udf.ImpalaUDA)
-        cur.release()
-
-        return result
+        with self._safe_raw_sql(statement) as cur:
+            return self._get_udfs(cur, udf.ImpalaUDA)
 
     def _get_udfs(self, cur, klass):
         def _to_type(x):
@@ -1189,7 +1188,7 @@ def compute_stats(
         cmd = f'COMPUTE {maybe_inc}STATS'
 
         stmt = self._table_command(cmd, name, database=database)
-        self.raw_sql(stmt)
+        self._safe_exec_sql(stmt)
 
     def invalidate_metadata(
         self,
@@ -1210,7 +1209,7 @@ def invalidate_metadata(
         stmt = 'INVALIDATE METADATA'
         if name is not None:
             stmt = self._table_command(stmt, name, database=database)
-        self.raw_sql(stmt)
+        self._safe_exec_sql(stmt)
 
     def refresh(self, name: str, database: str | None = None) -> None:
         """Reload HDFS block location metadata for a table.
@@ -1229,7 +1228,7 @@ def refresh(self, name: str, database: str | None = None) -> None:
         """
         # TODO(wesm): can this statement be cancelled?
         stmt = self._table_command('REFRESH', name, database=database)
-        self.raw_sql(stmt)
+        self._safe_exec_sql(stmt)
 
     def describe_formatted(
         self,
@@ -1292,7 +1291,8 @@ def column_stats(self, name, database=None):
         return self._exec_statement(stmt)
 
     def _exec_statement(self, stmt):
-        return self.fetch_from_cursor(self.raw_sql(stmt), schema=None)
+        with self._safe_raw_sql(stmt) as cur:
+            return self.fetch_from_cursor(cur, schema=None)
 
     def _table_command(self, cmd, name, database=None):
         qualified_name = self._fully_qualified_name(name, database)
diff --git a/ibis/backends/impala/client.py b/ibis/backends/impala/client.py
--- a/ibis/backends/impala/client.py
+++ b/ibis/backends/impala/client.py
@@ -305,7 +305,8 @@ def insert(
                 partition_schema=partition_schema,
                 overwrite=overwrite,
             )
-            return self._client.raw_sql(statement.compile())
+            self._client._safe_exec_sql(statement.compile())
+            return self
 
     def load_data(self, path, overwrite=False, partition=None):
         """Load data into an Impala table.
@@ -333,7 +334,8 @@ def load_data(self, path, overwrite=False, partition=None):
             overwrite=overwrite,
         )
 
-        return self._client.raw_sql(stmt.compile())
+        self._client._safe_exec_sql(stmt.compile())
+        return self
 
     @property
     def name(self):
@@ -348,7 +350,7 @@ def rename(self, new_name, database=None):
         if not m and database is None:
             database = self._database
         statement = RenameTable(self._qualified_name, new_name, new_database=database)
-        self._client.raw_sql(statement)
+        self._client._safe_exec_sql(statement)
 
         op = self.op().copy(name=statement.new_qualified_name)
         return type(self)(op)
@@ -385,7 +387,8 @@ def add_partition(self, spec, location=None):
         stmt = ddl.AddPartition(
             self._qualified_name, spec, part_schema, location=location
         )
-        return self._client.raw_sql(stmt)
+        self._client._safe_exec_sql(stmt)
+        return self
 
     def alter(
         self,
@@ -410,7 +413,8 @@ def alter(
 
         def _run_ddl(**kwds):
             stmt = AlterTable(self._qualified_name, **kwds)
-            return self._client.raw_sql(stmt)
+            self._client._safe_exec_sql(stmt)
+            return self
 
         return self._alter_table_helper(
             _run_ddl,
@@ -451,7 +455,8 @@ def alter_partition(
 
         def _run_ddl(**kwds):
             stmt = ddl.AlterPartition(self._qualified_name, spec, part_schema, **kwds)
-            return self._client.raw_sql(stmt)
+            self._client._safe_exec_sql(stmt)
+            return self
 
         return self._alter_table_helper(
             _run_ddl,
@@ -474,7 +479,8 @@ def drop_partition(self, spec):
         """Drop an existing table partition."""
         part_schema = self.partition_schema()
         stmt = ddl.DropPartition(self._qualified_name, spec, part_schema)
-        return self._client.raw_sql(stmt)
+        self._client._safe_exec_sql(stmt)
+        return self
 
     def partitions(self):
         """Return information about the table's partitions.


Test patch:
diff --git a/ibis/backends/impala/tests/conftest.py b/ibis/backends/impala/tests/conftest.py
--- a/ibis/backends/impala/tests/conftest.py
+++ b/ibis/backends/impala/tests/conftest.py
@@ -395,7 +395,7 @@ def mockcon():
 @pytest.fixture(scope="module")
 def kudu_table(con, test_data_db):
     name = 'kudu_backed_table'
-    con.raw_sql(
+    cur = con.raw_sql(
         f"""\
 CREATE TABLE {test_data_db}.{name} (
   a STRING,
@@ -408,6 +408,7 @@ def kudu_table(con, test_data_db):
   'kudu.num_tablet_replicas' = '1'
 )"""
     )
+    cur.close()
     yield con.table(name)
     con.drop_table(name, database=test_data_db)
 
diff --git a/ibis/backends/impala/tests/test_client.py b/ibis/backends/impala/tests/test_client.py
--- a/ibis/backends/impala/tests/test_client.py
+++ b/ibis/backends/impala/tests/test_client.py
@@ -1,4 +1,5 @@
 import datetime
+from contextlib import closing
 from posixpath import join as pjoin
 
 import pandas as pd
@@ -25,12 +26,6 @@ def test_hdfs_connect_function_is_public():
     assert hasattr(ibis.impala, "hdfs_connect")
 
 
-def test_cursor_garbage_collection(con):
-    for _ in range(5):
-        con.raw_sql('select 1').fetchall()
-        con.raw_sql('select 1').fetchone()
-
-
 def test_raise_ibis_error_no_hdfs(con_no_hdfs):
     # GH299
     with pytest.raises(com.IbisError):
@@ -62,9 +57,8 @@ def test_sql_with_limit(con):
 
 def test_raw_sql(con):
     query = 'SELECT * from functional_alltypes limit 10'
-    cur = con.raw_sql(query)
-    rows = cur.fetchall()
-    cur.release()
+    with closing(con.raw_sql(query)) as cur:
+        rows = cur.fetchall()
     assert len(rows) == 10
 
 
@@ -295,7 +289,7 @@ def test_datetime_to_int_cast(con):
 
 def test_set_option_with_dot(con):
     con.set_options({'request_pool': 'baz.quux'})
-    result = dict(row[:2] for row in con.raw_sql('set').fetchall())
+    result = con.get_options()
     assert result['REQUEST_POOL'] == 'baz.quux'
 
 
diff --git a/ibis/backends/impala/tests/test_ddl.py b/ibis/backends/impala/tests/test_ddl.py
--- a/ibis/backends/impala/tests/test_ddl.py
+++ b/ibis/backends/impala/tests/test_ddl.py
@@ -1,4 +1,5 @@
 import os
+from contextlib import closing
 from posixpath import join as pjoin
 
 import pytest
@@ -346,13 +347,16 @@ def test_query_delimited_file_directory(con, test_data_dir, tmp_db):
 @pytest.fixture
 def temp_char_table(con):
     name = "testing_varchar_support"
-    con.raw_sql(
-        f"""\
+    with closing(
+        con.raw_sql(
+            f"""\
 CREATE TABLE IF NOT EXISTS {name} (
   group1 VARCHAR(10),
   group2 CHAR(10)
 )"""
-    )
+        )
+    ):
+        pass
     assert name in con.list_tables(), name
     yield con.table(name)
     con.drop_table(name, force=True)