Source code for sqltrack.pandas

from __future__ import annotations

from typing import Union

from pandas import DataFrame

from .client import Client
from .util import SQL


def query_dataframe(client: Client, query: Union[str, SQL], parameters=()) -> DataFrame:
    """
    Run a query and return the result as a Pandas DataFrame.

    Parameters:
        client: The DB client to use
        query: The query to retrieve data
        parameters: Optional set of parameters passed to the cursor
    """
    with client as conn, conn.execute(query, parameters) as cursor:
        return DataFrame(cursor, columns=[col.name for col in cursor.description])