diff --git a/src/cursor.cpp b/src/cursor.cpp index 08002453..933c5359 100644 --- a/src/cursor.cpp +++ b/src/cursor.cpp @@ -338,6 +338,9 @@ static bool free_results(Cursor* self, int flags) if (self->colinfos) { + size_t n = (size_t)PyTuple_GET_SIZE(self->description); + for (size_t i = 0; i < n; ++i) + Py_XDECREF(self->colinfos[i].converter); PyMem_Free(self->colinfos); self->colinfos = 0; } @@ -495,6 +498,7 @@ bool InitColumnInfo(Cursor* cursor, SQLUSMALLINT iCol, ColumnInfo* pinfo) pinfo->sql_type = DataType; pinfo->column_size = ColumnSize; + pinfo->converter = 0; if (cursor->cnxn->hdbc == SQL_NULL_HANDLE) { @@ -1134,7 +1138,7 @@ static PyObject* Cursor_setinputsizes(PyObject* self, PyObject* sizes) PyErr_SetString(ProgrammingError, "Invalid cursor object."); return 0; } - + Cursor *cur = (Cursor*)self; if (Py_None == sizes) { @@ -2382,6 +2386,71 @@ static PyObject* Cursor_exit(PyObject* self, PyObject* args) Py_RETURN_NONE; } +static char setconv_doc[] = + "set_column_converter(column, converter) -> None\n\n" + "Register an output converter function or method that will be called\n" + "for the specified column of the results set created by the statement\n" + "most recently executed by the Cursor.\n\n" + "After executing a batch containing multiple statements, be sure to\n" + "use this function after calling nextset() if the results set for\n" + "which you need to register a converter is not the first statement\n" + "in the batch.\n\n" + "If no results set is active a RuntimeError exception will be raised.\n\n" + "column\n" + " The zero-based integer index of the column for which the converter\n" + " is to be registered. An IndexError exception will be raised if the\n" + " value is lower than zero or greater than or equal to the number of\n" + " columns in the results set.\n\n" + "converter\n" + " The converter function or method which will be called with a single\n" + " argument for the raw value retrieved from the database, and should\n" + " return the appropriately converted value. If the database value is\n" + " NULL, the argument given to the converter will be None. Otherwise\n" + " it will be a bytes object. If converter is None any existing con-\n" + " verter registered for the column is removed. If converter is not a\n" + " callable object a TypeError exception will be raised." + ; +static PyObject* Cursor_setconv(PyObject* self, PyObject* args) +{ + // Make sure we have a results set. + Cursor *cur = (Cursor*)self; + if (!PySequence_Check(cur->description)) + { + PyErr_SetString(PyExc_RuntimeError, "no results set is active"); + return 0; + } + + // Extract the function arguments. + int pos; + PyObject* func; + if (!PyArg_ParseTuple(args, "iO", &pos, &func)) + return 0; + + // Check the converter's type. + if (!PyCallable_Check(func) && func != Py_None) + { + PyErr_SetString(PyExc_TypeError, "converter is not callable"); + return 0; + } + + // Make sure this converter has a home. + int count = (int)PyTuple_GET_SIZE(cur->description); + if (pos < 0 || pos >= count) + { + PyErr_Format(PyExc_IndexError, "index %d out of range", pos); + return 0; + } + + // Free up any previous converter and install the new value. + Py_XDECREF(cur->colinfos[pos].converter); + if (func == Py_None) + cur->colinfos[pos].converter = NULL; + else + cur->colinfos[pos].converter = func; + Py_XINCREF(cur->colinfos[pos].converter); + + Py_RETURN_NONE; +} static PyMethodDef Cursor_methods[] = { @@ -2408,6 +2477,7 @@ static PyMethodDef Cursor_methods[] = { "skip", (PyCFunction)Cursor_skip, METH_VARARGS, skip_doc }, { "commit", (PyCFunction)Cursor_commit, METH_NOARGS, commit_doc }, { "rollback", (PyCFunction)Cursor_rollback, METH_NOARGS, rollback_doc }, + { "set_column_converter", (PyCFunction)Cursor_setconv, METH_VARARGS, setconv_doc }, {"cancel", (PyCFunction)Cursor_cancel, METH_NOARGS, cancel_doc}, {"__enter__", Cursor_enter, METH_NOARGS, enter_doc }, {"__exit__", Cursor_exit, METH_VARARGS, exit_doc }, diff --git a/src/cursor.h b/src/cursor.h index 657eb483..ad09aa05 100644 --- a/src/cursor.h +++ b/src/cursor.h @@ -31,6 +31,10 @@ struct ColumnInfo // of the integer types are the same size whether signed and unsigned, so we can allocate memory ahead of time // without knowing this. We use this during the fetch when converting to a Python integer or long. bool is_unsigned; + + // If not NULL, this represents a user-supplied function registered to perform custom + // transformation of the column's values. + PyObject* converter; }; struct ParamInfo @@ -117,10 +121,10 @@ struct Cursor // Parameter set array (used with executemany) unsigned char *paramArray; - + // Whether to use fast executemany with parameter arrays and other optimisations char fastexecmany; - + // The list of information for setinputsizes(). PyObject *inputsizes; diff --git a/src/getdata.cpp b/src/getdata.cpp index c401a35a..a153087b 100644 --- a/src/getdata.cpp +++ b/src/getdata.cpp @@ -692,7 +692,9 @@ PyObject* GetData(Cursor* cur, Py_ssize_t iCol) ColumnInfo* pinfo = &cur->colinfos[iCol]; // First see if there is a user-defined conversion. - + if (pinfo->converter) { + return GetDataUser(cur, iCol, pinfo->converter); + } if (cur->cnxn->map_sqltype_to_converter) { PyObject* func = Connection_GetConverter(cur->cnxn, pinfo->sql_type); if (func) { diff --git a/src/pyodbc.pyi b/src/pyodbc.pyi index 685cebf1..0a20ba59 100644 --- a/src/pyodbc.pyi +++ b/src/pyodbc.pyi @@ -646,6 +646,15 @@ class Cursor: """ ... + def set_column_converter(self, column: int, converter: Callable | None, /) -> None: + """Register an output converter for a column's value. + + Args: + column: The zero-based integer index of the column + converter: function/method to be registered (None to remove) + """ + ... + def fetchone(self) -> Row | None: """Retrieve the next row in the current result set for the query. diff --git a/tests/mysql_test.py b/tests/mysql_test.py index 37ecedaa..ba6fa47a 100644 --- a/tests/mysql_test.py +++ b/tests/mysql_test.py @@ -455,6 +455,96 @@ def test_emoticons_as_literal(cursor: pyodbc.Cursor): assert result == v +def test_column_converters(): + """Test the behavior of column-specific converters.""" + + # Create some converters. + def ucconv(raw): return raw.decode().upper() + class Converters: + now = datetime.now() + def __init__(self, name): self.name = name + def add_name(self, raw): return f"{raw.decode()} {self.name}" + @staticmethod + def show_raw(raw): return raw + @classmethod + def dt(cls, raw): return cls.now if b"e" in raw else None + object1 = Converters("Smith") + object2 = Converters("Čermák") + + # Create & populate a table with the test data. + names = "Leoš", "Kathy", "Renée", "Abdul", "George" + values = [[v] for v in names] + conn = connect() + cursor1 = conn.cursor() + cursor2 = conn.cursor() + cursor1.execute("drop table if exists t1") + conn.commit() + cursor1.execute("create table t1(v varchar(50))") + cursor1.executemany("insert into t1 values (?)", values) + + # Confirm that the documented exceptions get raised. + with pytest.raises(RuntimeError): + cursor1.set_column_converter(0, ucconv) + cursor1.execute("select v, v, v, v, v from t1") + cursor2.execute("select v, v, v, v, v from t1") + with pytest.raises(IndexError): + cursor1.set_column_converter(-1, ucconv) + with pytest.raises(IndexError): + cursor1.set_column_converter(5, ucconv) + with pytest.raises(TypeError): + cursor1.set_column_converter(0, "not callable") + + # Register initial conversions. + cursor1.set_column_converter(0, Converters.dt) + cursor1.set_column_converter(1, ucconv) + cursor1.set_column_converter(2, object1.add_name) + cursor1.set_column_converter(3, Converters.show_raw) + cursor2.set_column_converter(0, ucconv) + cursor2.set_column_converter(1, Converters.dt) + cursor2.set_column_converter(3, Converters.show_raw) + cursor2.set_column_converter(4, object2.add_name) + + # Create the validation tests. + expected_values = { + "cursor1": ( + (Converters.now, "LEOŠ", "Leoš Smith", "Leoš".encode(), "Leoš"), + ("Kathy", "KATHY", "Kathy Smith", b"Kathy", "Kathy"), + ("Renée", "RENÉE", "Renée Smith", "Renée".encode(), "Renée"), + ("Abdul", None, "Abdul Smith", b"Abdul", "Abdul"), + ("George", "George", "George", "George", "George"), + ), + "cursor2": ( + ("LEOŠ", Converters.now, "Leoš", "Leoš".encode(), "Leoš Čermák"), + ("KATHY", None, "Kathy", b"Kathy", "Kathy Čermák"), + ("RENÉE", Converters.now, "Renée".encode(), "Renée", "Renée Čermák"), + ("ABDUL", "ABDUL", b"Abdul", "Abdul", "Abdul Čermák"), + ("GEORGE", "GEORGE", b"George", "George", "George Čermák"), + ), + } + def check_row(cursor_name, index, row): + expected = expected_values[cursor_name][index] + assert tuple(row) == expected + + # Fetch a row at a time interleaving the cursors and modifying converter registrations as we go. + check_row("cursor1", 0, cursor1.fetchone()) + check_row("cursor2", 0, cursor2.fetchone()) + cursor1.set_column_converter(0, None) + check_row("cursor1", 1, cursor1.fetchone()) + check_row("cursor2", 1, cursor2.fetchone()) + cursor2.set_column_converter(2, Converters.show_raw) + cursor2.set_column_converter(3, None) + check_row("cursor1", 2, cursor1.fetchone()) + check_row("cursor2", 2, cursor2.fetchone()) + cursor1.set_column_converter(1, Converters.dt) + cursor2.set_column_converter(1, ucconv) + check_row("cursor1", 3, cursor1.fetchone()) + check_row("cursor2", 3, cursor2.fetchone()) + for i in range(len(cursor1.description)): + cursor1.set_column_converter(i, None) + check_row("cursor1", 4, cursor1.fetchone()) + check_row("cursor2", 4, cursor2.fetchone()) + + @lru_cache def _generate_str(length, encoding=None): """ diff --git a/tests/postgresql_test.py b/tests/postgresql_test.py index 96c794b5..639e2d3b 100644 --- a/tests/postgresql_test.py +++ b/tests/postgresql_test.py @@ -4,6 +4,7 @@ # -*- coding: utf-8 -*- import os, uuid +from datetime import datetime from decimal import Decimal from typing import Iterator @@ -595,6 +596,96 @@ def convert(value): assert value == '123.45' +def test_column_converters(): + """Test the behavior of column-specific converters.""" + + # Create some converters. + def ucconv(raw): return raw.decode().upper() + class Converters: + now = datetime.now() + def __init__(self, name): self.name = name + def add_name(self, raw): return f"{raw.decode()} {self.name}" + @staticmethod + def show_raw(raw): return raw + @classmethod + def dt(cls, raw): return cls.now if b"e" in raw else None + object1 = Converters("Smith") + object2 = Converters("Czerny") + + # Create & populate a table with the test values. + names = "Leon", "Kathy", "Yentl", "Abdul", "George" + values = [[v] for v in names] + conn = connect() + cursor1 = conn.cursor() + cursor2 = conn.cursor() + cursor1.execute("drop table if exists t1") + conn.commit() + cursor1.execute("create table t1(v varchar(50))") + cursor1.executemany("insert into t1 values (?)", values) + + # Confirm that the documented exceptions get raised. + with pytest.raises(RuntimeError): + cursor1.set_column_converter(0, ucconv) + cursor1.execute("select v, v, v, v, v from t1") + cursor2.execute("select v, v, v, v, v from t1") + with pytest.raises(IndexError): + cursor1.set_column_converter(-1, ucconv) + with pytest.raises(IndexError): + cursor1.set_column_converter(5, ucconv) + with pytest.raises(TypeError): + cursor1.set_column_converter(0, "not callable") + + # Register initial conversions. + cursor1.set_column_converter(0, Converters.dt) + cursor1.set_column_converter(1, ucconv) + cursor1.set_column_converter(2, object1.add_name) + cursor1.set_column_converter(3, Converters.show_raw) + cursor2.set_column_converter(0, ucconv) + cursor2.set_column_converter(1, Converters.dt) + cursor2.set_column_converter(3, Converters.show_raw) + cursor2.set_column_converter(4, object2.add_name) + + # Create the validation tests. + expected_values = { + "cursor1": ( + (Converters.now, "LEON", "Leon Smith", b"Leon", "Leon"), + ("Kathy", "KATHY", "Kathy Smith", b"Kathy", "Kathy"), + ("Yentl", "YENTL", "Yentl Smith", b"Yentl", "Yentl"), + ("Abdul", None, "Abdul Smith", b"Abdul", "Abdul"), + ("George", "George", "George", "George", "George"), + ), + "cursor2": ( + ("LEON", Converters.now, "Leon", b"Leon", "Leon Czerny"), + ("KATHY", None, "Kathy", b"Kathy", "Kathy Czerny"), + ("YENTL", Converters.now, b"Yentl", "Yentl", "Yentl Czerny"), + ("ABDUL", "ABDUL", b"Abdul", "Abdul", "Abdul Czerny"), + ("GEORGE", "GEORGE", b"George", "George", "George Czerny"), + ), + } + def check_row(cursor_name, index, row): + expected = expected_values[cursor_name][index] + assert tuple(row) == expected + + # Fetch a row at a time interleaving the cursors and modifying converter registrations as we go. + check_row("cursor1", 0, cursor1.fetchone()) + check_row("cursor2", 0, cursor2.fetchone()) + cursor1.set_column_converter(0, None) + check_row("cursor1", 1, cursor1.fetchone()) + check_row("cursor2", 1, cursor2.fetchone()) + cursor2.set_column_converter(2, Converters.show_raw) + cursor2.set_column_converter(3, None) + check_row("cursor1", 2, cursor1.fetchone()) + check_row("cursor2", 2, cursor2.fetchone()) + cursor1.set_column_converter(1, Converters.dt) + cursor2.set_column_converter(1, ucconv) + check_row("cursor1", 3, cursor1.fetchone()) + check_row("cursor2", 3, cursor2.fetchone()) + for i in range(len(cursor1.description)): + cursor1.set_column_converter(i, None) + check_row("cursor1", 4, cursor1.fetchone()) + check_row("cursor2", 4, cursor2.fetchone()) + + def test_refcount_encoding(): """ Ensure we handle the reference count to `encoding` properly. In the past we freed a diff --git a/tests/sqlite_test.py b/tests/sqlite_test.py index 6d4150ef..4c5be49f 100644 --- a/tests/sqlite_test.py +++ b/tests/sqlite_test.py @@ -726,3 +726,92 @@ def test_pickling(cnxn: pyodbc.Connection): pickled_rows = pickle.dumps(original_rows) unpickled_rows = pickle.loads(pickled_rows) assert unpickled_rows == original_rows + + +def test_column_converters(cnxn: pyodbc.Connection): + """Test the behavior of column-specific converters.""" + + # Create some converters. + def ucconv(raw): return raw.decode().upper() + class Converters: + now = datetime.now() + def __init__(self, name): self.name = name + def add_name(self, raw): return f"{raw.decode()} {self.name}" + @staticmethod + def show_raw(raw): return raw + @classmethod + def dt(cls, raw): return cls.now if b"e" in raw else None + object1 = Converters("Smith") + object2 = Converters("Čermák") + + # Create & populate a table with the test values. + names = "Leoš", "Kathy", "Renée", "Abdul", "George" + values = [[v] for v in names] + cursor1 = cnxn.cursor() + cursor2 = cnxn.cursor() + cursor1.execute("drop table if exists t1") + cnxn.commit() + cursor1.execute("create table t1(v varchar(50))") + cursor1.executemany("insert into t1 values (?)", values) + + # Confirm that the documented exceptions get raised. + with pytest.raises(RuntimeError): + cursor1.set_column_converter(0, ucconv) + cursor1.execute("select v, v, v, v, v from t1") + cursor2.execute("select v, v, v, v, v from t1") + with pytest.raises(IndexError): + cursor1.set_column_converter(-1, ucconv) + with pytest.raises(IndexError): + cursor1.set_column_converter(5, ucconv) + with pytest.raises(TypeError): + cursor1.set_column_converter(0, "not callable") + + # Register initial conversions. + cursor1.set_column_converter(0, Converters.dt) + cursor1.set_column_converter(1, ucconv) + cursor1.set_column_converter(2, object1.add_name) + cursor1.set_column_converter(3, Converters.show_raw) + cursor2.set_column_converter(0, ucconv) + cursor2.set_column_converter(1, Converters.dt) + cursor2.set_column_converter(3, Converters.show_raw) + cursor2.set_column_converter(4, object2.add_name) + + # Create the validation tests. + expected_values = { + "cursor1": ( + (Converters.now, "LEOŠ", "Leoš Smith", "Leoš".encode(), "Leoš"), + ("Kathy", "KATHY", "Kathy Smith", b"Kathy", "Kathy"), + ("Renée", "RENÉE", "Renée Smith", "Renée".encode(), "Renée"), + ("Abdul", None, "Abdul Smith", b"Abdul", "Abdul"), + ("George", "George", "George", "George", "George"), + ), + "cursor2": ( + ("LEOŠ", Converters.now, "Leoš", "Leoš".encode(), "Leoš Čermák"), + ("KATHY", None, "Kathy", b"Kathy", "Kathy Čermák"), + ("RENÉE", Converters.now, "Renée".encode(), "Renée", "Renée Čermák"), + ("ABDUL", "ABDUL", b"Abdul", "Abdul", "Abdul Čermák"), + ("GEORGE", "GEORGE", b"George", "George", "George Čermák"), + ), + } + def check_row(cursor_name, index, row): + expected = expected_values[cursor_name][index] + assert tuple(row) == expected + + # Fetch a row at a time interleaving the cursors and modifying converter registrations as we go. + check_row("cursor1", 0, cursor1.fetchone()) + check_row("cursor2", 0, cursor2.fetchone()) + cursor1.set_column_converter(0, None) + check_row("cursor1", 1, cursor1.fetchone()) + check_row("cursor2", 1, cursor2.fetchone()) + cursor2.set_column_converter(2, Converters.show_raw) + cursor2.set_column_converter(3, None) + check_row("cursor1", 2, cursor1.fetchone()) + check_row("cursor2", 2, cursor2.fetchone()) + cursor1.set_column_converter(1, Converters.dt) + cursor2.set_column_converter(1, ucconv) + check_row("cursor1", 3, cursor1.fetchone()) + check_row("cursor2", 3, cursor2.fetchone()) + for i in range(len(cursor1.description)): + cursor1.set_column_converter(i, None) + check_row("cursor1", 4, cursor1.fetchone()) + check_row("cursor2", 4, cursor2.fetchone()) diff --git a/tests/sqlserver_test.py b/tests/sqlserver_test.py index 6feee9ec..c61a5ec1 100755 --- a/tests/sqlserver_test.py +++ b/tests/sqlserver_test.py @@ -1277,6 +1277,98 @@ def convert2(value): assert value == '123.45' +def test_column_converters(): + """Test the behavior of column-specific converters.""" + + # Create some converters. + legacy = "11" in DRIVER or "13" in DRIVER + encoding = "UTF-16LE" if legacy else "utf-8" + def ucconv(raw): return raw.decode(encoding).upper() + class Converters: + now = datetime.now() + def __init__(self, name): self.name = name + def add_name(self, raw): return f"{raw.decode(encoding)} {self.name}" + @staticmethod + def show_raw(raw): return raw.decode(encoding).encode("utf-8") if legacy else raw + @classmethod + def dt(cls, raw): return cls.now if b"e" in raw else None + object1 = Converters("Smith") + object2 = Converters("Čermák") + + # Create & populate a table with a known encoding, with multiple active results sets enabled. + names = "Leoš", "Kathy", "Renée", "Abdul", "George" + values = [[v] for v in names] + conn = pyodbc.connect(f"MARS_Connection=Yes;{CNXNSTR}") + cursor1 = conn.cursor() + cursor2 = conn.cursor() + cursor1.execute("drop table if exists t1") + conn.commit() + cursor1.execute("create table t1(v varchar(50) collate Latin1_General_100_CI_AS_SC_UTF8)") + cursor1.executemany("insert into t1 values (?)", values) + + # Confirm that the documented exceptions get raised. + with pytest.raises(RuntimeError): + cursor1.set_column_converter(0, ucconv) + cursor1.execute("select v, v, v, v, v from t1") + cursor2.execute("select v, v, v, v, v from t1") + with pytest.raises(IndexError): + cursor1.set_column_converter(-1, ucconv) + with pytest.raises(IndexError): + cursor1.set_column_converter(5, ucconv) + with pytest.raises(TypeError): + cursor1.set_column_converter(0, "not callable") + + # Register initial conversions. + cursor1.set_column_converter(0, Converters.dt) + cursor1.set_column_converter(1, ucconv) + cursor1.set_column_converter(2, object1.add_name) + cursor1.set_column_converter(3, Converters.show_raw) + cursor2.set_column_converter(0, ucconv) + cursor2.set_column_converter(1, Converters.dt) + cursor2.set_column_converter(3, Converters.show_raw) + cursor2.set_column_converter(4, object2.add_name) + + # Create the validation tests. + expected_values = { + "cursor1": ( + (Converters.now, "LEOŠ", "Leoš Smith", "Leoš".encode(), "Leoš"), + ("Kathy", "KATHY", "Kathy Smith", b"Kathy", "Kathy"), + ("Renée", "RENÉE", "Renée Smith", "Renée".encode(), "Renée"), + ("Abdul", None, "Abdul Smith", b"Abdul", "Abdul"), + ("George", "George", "George", "George", "George"), + ), + "cursor2": ( + ("LEOŠ", Converters.now, "Leoš", "Leoš".encode(), "Leoš Čermák"), + ("KATHY", None, "Kathy", b"Kathy", "Kathy Čermák"), + ("RENÉE", Converters.now, "Renée".encode(), "Renée", "Renée Čermák"), + ("ABDUL", "ABDUL", b"Abdul", "Abdul", "Abdul Čermák"), + ("GEORGE", "GEORGE", b"George", "George", "George Čermák"), + ), + } + def check_row(cursor_name, index, row): + expected = expected_values[cursor_name][index] + assert tuple(row) == expected + + # Fetch a row at a time interleaving the cursors and modifying converter registrations as we go. + check_row("cursor1", 0, cursor1.fetchone()) + check_row("cursor2", 0, cursor2.fetchone()) + cursor1.set_column_converter(0, None) + check_row("cursor1", 1, cursor1.fetchone()) + check_row("cursor2", 1, cursor2.fetchone()) + cursor2.set_column_converter(2, Converters.show_raw) + cursor2.set_column_converter(3, None) + check_row("cursor1", 2, cursor1.fetchone()) + check_row("cursor2", 2, cursor2.fetchone()) + cursor1.set_column_converter(1, Converters.dt) + cursor2.set_column_converter(1, ucconv) + check_row("cursor1", 3, cursor1.fetchone()) + check_row("cursor2", 3, cursor2.fetchone()) + for i in range(len(cursor1.description)): + cursor1.set_column_converter(i, None) + check_row("cursor1", 4, cursor1.fetchone()) + check_row("cursor2", 4, cursor2.fetchone()) + + def test_too_large(cursor: pyodbc.Cursor): """Ensure error raised if insert fails due to truncation""" value = 'x' * 1000