This is an automated email from the ASF dual-hosted git repository.
lidavidm pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-adbc.git
The following commit(s) were added to refs/heads/main by this push:
new c324fae0e feat(python/adbc_driver_manager): add cursor() arg to set
options (#2589)
c324fae0e is described below
commit c324fae0e22180aaa79278c29f0f2be988695181
Author: David Li <[email protected]>
AuthorDate: Mon Mar 10 18:37:51 2025 -0400
feat(python/adbc_driver_manager): add cursor() arg to set options (#2589)
- Add a keyword-only argument to set statement options directly when
creating a cursor, which makes it a little clearer what's going on.
- Allow using True and False directly as option values (they will get
converted to ENABLED and DISABLED, respectively).
- Add the `use_copy` option to the PostgreSQL options enum.
- Add an example of this with PostgreSQL.
Fixes #2093.
Fixes #2146.
---
.../python/recipe/postgresql_execute_nocopy.py | 74 ++++++++++++++++++++++
.../recipe/postgresql_execute_nocopy.py.stdout.txt | 6 ++
.../adbc_driver_manager/_lib.pyx | 27 ++++++++
.../adbc_driver_manager/dbapi.py | 26 ++++++--
.../adbc_driver_postgresql/__init__.py | 6 ++
5 files changed, 135 insertions(+), 4 deletions(-)
diff --git a/docs/source/python/recipe/postgresql_execute_nocopy.py
b/docs/source/python/recipe/postgresql_execute_nocopy.py
new file mode 100644
index 000000000..b9f63bcda
--- /dev/null
+++ b/docs/source/python/recipe/postgresql_execute_nocopy.py
@@ -0,0 +1,74 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+# RECIPE CATEGORY: PostgreSQL
+# RECIPE KEYWORDS: statement options
+# RECIPE STARTS HERE
+
+#: PostgreSQL does not support ``COPY`` for all kinds of queries, for example,
+#: queries that use ``SHOW``. But the ADBC driver tries to execute queries
+#: with COPY by default since it is faster for large result sets. In this
+#: case, you can explicitly disable the ``COPY`` optimization.
+
+import os
+
+import adbc_driver_postgresql.dbapi
+
+uri = os.environ["ADBC_POSTGRESQL_TEST_URI"]
+conn = adbc_driver_postgresql.dbapi.connect(uri)
+
+#: The option can be set when creating the cursor:
+
+with conn.cursor(
+ adbc_stmt_kwargs={
+ adbc_driver_postgresql.StatementOptions.USE_COPY.value: False,
+ }
+) as cur:
+ cur.execute("SHOW ALL")
+ print(cur.fetch_arrow_table().schema)
+ # Output:
+ # name: string
+ # setting: string
+ # description: string
+
+#: Or it can be set afterwards:
+
+with conn.cursor() as cur:
+ cur.adbc_statement.set_options(
+ **{
+ adbc_driver_postgresql.StatementOptions.USE_COPY.value: False,
+ }
+ )
+ cur.execute("SHOW ALL")
+ print(cur.fetch_arrow_table().schema)
+ # Output:
+ # name: string
+ # setting: string
+ # description: string
+
+#: Without the option, the query fails as the driver attempts to execute the
+#: query with ``COPY``:
+
+with conn.cursor() as cur:
+ try:
+ cur.execute("SHOW ALL")
+ except conn.Error:
+ pass
+ else:
+ raise RuntimeError("Expected error")
+
+conn.close()
diff --git a/docs/source/python/recipe/postgresql_execute_nocopy.py.stdout.txt
b/docs/source/python/recipe/postgresql_execute_nocopy.py.stdout.txt
new file mode 100644
index 000000000..0ac1231ec
--- /dev/null
+++ b/docs/source/python/recipe/postgresql_execute_nocopy.py.stdout.txt
@@ -0,0 +1,6 @@
+name: string
+setting: string
+description: string
+name: string
+setting: string
+description: string
diff --git a/python/adbc_driver_manager/adbc_driver_manager/_lib.pyx
b/python/adbc_driver_manager/adbc_driver_manager/_lib.pyx
index 9f4b80df3..21afe9d3c 100644
--- a/python/adbc_driver_manager/adbc_driver_manager/_lib.pyx
+++ b/python/adbc_driver_manager/adbc_driver_manager/_lib.pyx
@@ -637,6 +637,15 @@ cdef class AdbcDatabase(_AdbcHandle):
c_value = value
status = AdbcDatabaseSetOption(
&self.database, c_key, c_value, &c_error)
+ elif isinstance(value, bool):
+ if value:
+ value = ADBC_OPTION_VALUE_ENABLED
+ else:
+ value = ADBC_OPTION_VALUE_DISABLED
+ value = _to_bytes(value, "option value")
+ c_value = value
+ status = AdbcDatabaseSetOption(
+ &self.database, c_key, c_value, &c_error)
elif isinstance(value, bytes):
c_value = value
status = AdbcDatabaseSetOptionBytes(
@@ -1021,6 +1030,15 @@ cdef class AdbcConnection(_AdbcHandle):
c_value = value
status = AdbcConnectionSetOption(
&self.connection, c_key, c_value, &c_error)
+ elif isinstance(value, bool):
+ if value:
+ value = ADBC_OPTION_VALUE_ENABLED
+ else:
+ value = ADBC_OPTION_VALUE_DISABLED
+ value = _to_bytes(value, "option value")
+ c_value = value
+ status = AdbcConnectionSetOption(
+ &self.connection, c_key, c_value, &c_error)
elif isinstance(value, bytes):
c_value = value
status = AdbcConnectionSetOptionBytes(
@@ -1467,6 +1485,15 @@ cdef class AdbcStatement(_AdbcHandle):
c_value = value
status = AdbcStatementSetOption(
&self.statement, c_key, c_value, &c_error)
+ elif isinstance(value, bool):
+ if value:
+ value = ADBC_OPTION_VALUE_ENABLED
+ else:
+ value = ADBC_OPTION_VALUE_DISABLED
+ value = _to_bytes(value, "option value")
+ c_value = value
+ status = AdbcStatementSetOption(
+ &self.statement, c_key, c_value, &c_error)
elif isinstance(value, bytes):
c_value = value
status = AdbcStatementSetOptionBytes(
diff --git a/python/adbc_driver_manager/adbc_driver_manager/dbapi.py
b/python/adbc_driver_manager/adbc_driver_manager/dbapi.py
index fea644514..9dfc4e551 100644
--- a/python/adbc_driver_manager/adbc_driver_manager/dbapi.py
+++ b/python/adbc_driver_manager/adbc_driver_manager/dbapi.py
@@ -338,9 +338,20 @@ class Connection(_Closeable):
if self._commit_supported:
self._conn.commit()
- def cursor(self) -> "Cursor":
- """Create a new cursor for querying the database."""
- return Cursor(self)
+ def cursor(
+ self,
+ *,
+ adbc_stmt_kwargs: Optional[Dict[str, Any]] = None,
+ ) -> "Cursor":
+ """
+ Create a new cursor for querying the database.
+
+ Parameters
+ ----------
+ adbc_stmt_kwargs : dict, optional
+ ADBC-specific options to pass to the underlying ADBC statement.
+ """
+ return Cursor(self, adbc_stmt_kwargs)
def rollback(self) -> None:
"""Explicitly rollback."""
@@ -571,7 +582,11 @@ class Cursor(_Closeable):
Do not create this object directly; use Connection.cursor().
"""
- def __init__(self, conn: Connection) -> None:
+ def __init__(
+ self,
+ conn: Connection,
+ adbc_stmt_kwargs: Optional[Dict[str, Any]] = None,
+ ) -> None:
# Must be at top in case __init__ is interrupted and then __del__ is
called
self._closed = True
self._conn = conn
@@ -583,6 +598,9 @@ class Cursor(_Closeable):
self._arraysize = 1
self._rowcount = -1
+ if adbc_stmt_kwargs:
+ self._stmt.set_options(**adbc_stmt_kwargs)
+
@property
def arraysize(self) -> int:
"""The number of rows to fetch at a time with fetchmany()."""
diff --git a/python/adbc_driver_postgresql/adbc_driver_postgresql/__init__.py
b/python/adbc_driver_postgresql/adbc_driver_postgresql/__init__.py
index 467794a78..34abece63 100644
--- a/python/adbc_driver_postgresql/adbc_driver_postgresql/__init__.py
+++ b/python/adbc_driver_postgresql/adbc_driver_postgresql/__init__.py
@@ -34,6 +34,12 @@ class StatementOptions(enum.Enum):
#: actual size may differ.
BATCH_SIZE_HINT_BYTES = "adbc.postgresql.batch_size_hint_bytes"
+ #: Enable or disable the ``COPY`` optimization (default: enabled).
+ #:
+ #: This is necessary for some queries since PostgreSQL does not support
+ #: ``COPY`` with those queries, e.g. queries using ``SHOW``.
+ USE_COPY = "adbc.postgresql.use_copy"
+
def connect(uri: str) -> adbc_driver_manager.AdbcDatabase:
"""Create a low level ADBC connection to PostgreSQL."""
