Retrieving Multiple Result-Sets from SQLAlchemy

SQLAlchemy is a great Python-based database client, but, traditionally, it leaves you stuck when it comes to stored-procedures that return more than one dataset. This means that you’d have to either call separate queries or merge multiple datasets into one large, unnatural one. However, there is a way to read multiple datasets but it requires accessing the raw MySQL layer (which isn’t too bad).

This is the test-routine:

delimiter //

CREATE PROCEDURE `get_sets`()
BEGIN
    SELECT
        'value1' `series1_col1`,
        'value2' `series1_col2`;

    SELECT
        'value3' `series2_col1`,
        'value4' `series2_col2`;

    SELECT
        'value5' `series3_col1`,
        'value6' `series3_col2`;
END//

delimiter ;

The code:

import json

import sqlalchemy.pool

def _run_query(connection, query, parameters={}):
    sets = []

    try:
        cursor = connection.cursor()

        cursor.execute(query, parameters)

        while 1:
            #(column_name, type_, ignore_, ignore_, ignore_, null_ok, column_flags)
            names = [c[0] for c in cursor.description]

            set_ = []
            while 1:
                row_raw = cursor.fetchone()
                if row_raw is None:
                    break

                row = dict(zip(names, row_raw))
                set_.append(row)

            sets.append(list(set_))

            if cursor.nextset() is None:
                break

            # nextset() doesn't seem to be sufficiant to tell the end.
            if cursor.description is None:
                break
    finally:
        # Return the connection to the pool (won't actually close).
        connection.close()

    return sets

def _pretty_json_dumps(data):
    return json.dumps(
            data,
            sort_keys=True,
            indent=4, 
            separators=(',', ': ')) + "\n"

def _main():
    dsn = 'mysql+mysqldb://root:root@localhost:3306/test_database'

    engine = sqlalchemy.create_engine(
                dsn, 
                pool_recycle=7200,
                poolclass=sqlalchemy.pool.NullPool)

    # Grab a raw connection from the connection-pool.
    connection = engine.raw_connection()

    query = 'CALL get_sets()'
    sets = _run_query(connection, query)

    print(_pretty_json_dumps(sets))

if __name__ == '__main__':
    _main()

The output:

[
    [
        {
            "series1_col1": "value1",
            "series1_col2": "value2"
        }
    ],
    [
        {
            "series2_col1": "value3",
            "series2_col2": "value4"
        }
    ],
    [
        {
            "series3_col1": "value5",
            "series3_col2": "value6"
        }
    ]
]

Things to observe in the example:

  • The query parameters are still escaped (our parameters have spaces in them), even though we have to use classic Python string-substitution formatting with the raw connection-objects.
  • It’s up to us to extract the column-names from the cursor for each dataset.
  • The resulting datasets can’t be captured as generators, as they have to be read entirely before jumping to the next dataset. Technically, you can yield each dataset, but this has almost no usefulness since you’d rarely be required need to read through them sequentially and you’d only benefit if there were a large number of datasets.
  • The raw_connection() method claims a connection from the pool, and its close() method will return it to the pool without actually closing it.
  • I added pool_recycle for good measure. This is an enormous pain to have to deal with, if you’re new to SA and your connections keep “going away” because MySQL is closing them before SA can recycle them.

REFERENCE: Multiple Result Sets

Advertisements