Handle pyodbc.Row returned by ODBCHook with SELECT statememt

[Joffreybvn]

Hello,

Using the ODBCHook with the SQLExecuteQueryOperator works fine for any SQL query that do not return data. However, it fails when doing a SELECT, because the ODBCHook (powered by PyODBC) returns a list of pyodbc.Row objects (instead of a list of tuple like psycopg2). More specifically, the XCom backend receive those Row objects and fails to serialize them.

Example:

SQLExecuteQueryOperator(
    task_id="execute_select_statement",
    conn_id="odbc_connection",  # Setup as "ODBC" in Airflow 'Connection' page
    sql="""SELECT "id", "name" from "schema.table" limit 10;""",
    do_xcom_push=True,
    dag=dag,
)

Will throw:

Traceback (most recent call last):
  File "/opt/app-root/lib64/python3.9/site-packages/airflow/utils/json.py", line 91, in default
    return serialize(o)
  File "/opt/app-root/lib64/python3.9/site-packages/airflow/serialization/serde.py", line 171, in serialize
    raise TypeError(f"cannot serialize object of type {cls}")
TypeError: cannot serialize object of type <class 'pyodbc.Row'>
During handling of the above exception, another exception occurred:
Traceback (most recent call last):
  File "/opt/app-root/lib64/python3.9/site-packages/airflow/utils/session.py", line 73, in wrapper
    return func(*args, **kwargs)
  File "/opt/app-root/lib64/python3.9/site-packages/airflow/models/taskinstance.py", line 2354, in xcom_push
    XCom.set(
  File "/opt/app-root/lib64/python3.9/site-packages/airflow/utils/session.py", line 73, in wrapper
    return func(*args, **kwargs)
  File "/opt/app-root/lib64/python3.9/site-packages/airflow/models/xcom.py", line 237, in set
    value = cls.serialize_value(
  File "/opt/app-root/lib64/python3.9/site-packages/airflow/models/xcom.py", line 632, in serialize_value
    return json.dumps(value, cls=XComEncoder).encode("UTF-8")
  File "/usr/lib64/python3.9/json/__init__.py", line 234, in dumps
    return cls(
  File "/opt/app-root/lib64/python3.9/site-packages/airflow/utils/json.py", line 104, in encode
    return super().encode(o)
  File "/usr/lib64/python3.9/json/encoder.py", line 199, in encode
    chunks = self.iterencode(o, _one_shot=True)
  File "/usr/lib64/python3.9/json/encoder.py", line 257, in iterencode
    return _iterencode(o, 0)
  File "/opt/app-root/lib64/python3.9/site-packages/airflow/utils/json.py", line 93, in default
    return super().default(o)
  File "/usr/lib64/python3.9/json/encoder.py", line 179, in default
    raise TypeError(f'Object of type {o.__class__.__name__} '
TypeError: Object of type Row is not JSON serializable

Solution

I would like to improve the behavior of the SQLExecuteQueryOperator so that result of a SELECT statement made with the ODBCHook can be passed via XCom and used by downstream tasks. Yes, I'm whilling to submit a PR ✅.

But, given that a pyodbc.Row object can be converted into a sequence (list / tuple), or into a dictionnary (with column' names as key), they are several solutions possible:

1. Creating two handlers: Like the fetch_all_handler method used as default handler of the SQLExecuteQueryOperator, I'd create two handlers methods (one to get result as list, one to get result as dict) in the ODBC provider. + some docs to inform that those methods exists.
2. Same as 1 + Create an ODBCOperator on top of SQLExecuteQueryOperator, which would have one of the new handler set as default.
3. Create an ODBCOperators and customize the _process_output method. With an extra parameter in the Operator for the user to define how he want to convert the Row object.
4. Adding support to Row in the XCom backend

And here I need the light of the community:

My prefered solution is the (2): Simple inheritance with minimal maintenance.

But, the _process_output method (solution 3), created to run just before a return / just before reaching the XCom backend, sounds also like a good place to fix this. Maybe the scope of the handler is not to prepare for serialization ?

What do you think ?

[potiuk]

I vote for option 5. Make the results returned by ODBCHook Json Serialzable. We've done something simillar for Databricks (but there the operator was already implemented) https://github.com/apache/airflow/pull/31780/files - I believe the solution is to make a very similar change to convert whatever PyODBC returns from the hook (non-serializable) into serializables.

[dabla]

I would personally not interfer with the result type returned by the OdbcHook, that would be more restrictive if someone wants to use it in a python method with the PythonOperator don't you think? To me a Hook is more a low level component in Airflow and Operators are more high level, thus I would do the conversion within the operator.

[potiuk]

Not really - depends what objects you return. if (similarly to the PR above) the Row objects returned follow the same semantics /fields as the original objects, I see no problem with that. The change above uses NamedTuple. What is the problem you think it is going to create @dabla ?

[dabla]

It would not be a real problem of course. I would expect the OdbcHook to return a pyodbc.row in the python method that's being invoked in an PythonOperator and if for example I would like to invoke the cursor_description method of the pyodbc.row class that would not work as it has been serialized. I would expect this behaviour when using an operator as then you don't want to deal with that low level stuff. This is of course a personal opinion, doesn't mean it's the right one ;)

[potiuk]

Yeah. Can be both - but there is a value in not having to create a separate operator just to handle different type. The more "standard" the hooks are the better - for example currently we will can handle column-lineage with upcoming open-lineage integration automatically when SQLOperator is used - independently what database we talk to. If we have a separate operator, we would have to add such integration for both ODBCSQLOperator and SQLOperator. So there is a value in having just one operator handling all SQL.

And also in the future it will be easier if the value returned by different Hooks will be pretty much the same - that for example will enable column-level lineage for @task decorated tasks which are using the Hook. So if only we can easily use a bit more "standard" approach there and serialize those values, that would be better IMHO.

d if for example I would like to invoke the cursor_description method of the pyodbc.row class that would not work as it has been serialized

I'd argue quite the opposite. You do NOT want to use row.description. Look at the API of DBApiHook. It already has different, standardized way of retrieving description which is better optimized. also supporting the case where you have many rows and multiple statements (descriptions array property and last_description() method).

[bolkedebruin]

For future reference, to make custom objects serializable either do the following:

1. Add a serialize / deserialize method to the object
2. Add a serializer/deserializer into airflow.serialization.serializers
3. Decorate the object with @DataClass or @attr

The is no need for custom serialization schemes.

[potiuk]

See the comment in #32319 (comment) @bolkedebruin . I think you are mixing things. The implementation there was an attempt to convert something returned by DBAPI (which unfortunately is not standardized by Python DB API specification) and implemented differently by different DBApi Implementations into something that we already standardized (the structure of serializable tuples that DBAPIHook returns).

I think:

1. is not applicable because we get those objects returned by the DBAPI cursor
2. Is not appicable because we want to make it available as output of th DBApiHook in standardized (airflow standard) serializable (to be used by Xcom) form
3. is also not applicable for the same reason as 1).

But maybe I am wrong and misunderstood your point :) .

[bolkedebruin]

I need to do more documentation / explanation around this as we discussed. serde serializes objects so that they can be used by XCom. Exactly, your number 2 :-). Standard XCom uses a second pass serializer in the form of JSON to be able to put it into the database.

Note: currently, serde is not part of the public API, but it does not require that to function. However a number 2 implementation requires an upgrade of Airflow .

The error shown here:

Traceback (most recent call last):
  File "/opt/app-root/lib64/python3.9/site-packages/airflow/utils/json.py", line 91, in default
    return serialize(o)
  File "/opt/app-root/lib64/python3.9/site-packages/airflow/serialization/serde.py", line 171, in serialize
    raise TypeError(f"cannot serialize object of type {cls}")
TypeError: cannot serialize object of type <class 'pyodbc.Row'>

could have been solved by creating a serializer in airflow.serialization.serializers like this:

# lazy loading for performance reasons
serializers = [
    "pyodbc.Row"
]
deserializers = serializers

__version__ = 1


if TYPE_CHECKING:
    from airflow.serialization.serde import U


def serialize(o: object) -> tuple[U, str, int, bool]:
    import pyodbc

    o = cast(pyodbc.Row, o)

    columns: list[tuple[str, type]] = [col[:2] for col in o.cursor_description]
    row_object = NamedTuple("Row", columns)  # type: ignore[misc]
    row = row_object(*o)

    return row, qualname(o), __version__, True


def deserialize(classname: str, version: int, data: object) -> Any:
    import pyodbc

    if version > __version__:
        raise TypeError("serialized version is newer than class version")

    if classname == qualname(pyodbc.Row):
        return data

    raise TypeError(f"do not know how to deserialize {classname}")

Now the linked PR has a broader scope and thus implements this differently.