Source code for sqltrack.queries

from __future__ import annotations

from typing import Union

from collections import namedtuple

from .client import Client
from .util import SQL


__all__ = [
    "first_row",
    "first_value",
    "first_values",
]


[docs]def first_row(client: Client, query: Union[str, SQL], parameters=()):
    """
    Execute a query and return the first matching row, if any.

    Parameters:
        cursor: Cursor to use
        query: Query to execute
        parameters: Optional parameters
    """
    with client.cursor() as cursor:
        cursor.execute(query, parameters)
        return cursor.fetchone()


[docs]def first_value(client: Client, query: Union[str, SQL], parameters=()):
    """
    Execute a query and return the first value of the first matching row, if any.

    Parameters:
        cursor: Cursor to use
        query: Query to execute
        parameters: Optional parameters
    """
    with client:
        row = first_row(client, query, parameters)
        if row is None:
            return None
        return row[0]


[docs]def first_values(client: Client, query: Union[str, SQL], parameters=()):
    """
    Execute a query and return the first value of each matching row.

    Parameters:
        cursor: Cursor to use
        query: Query to execute
        parameters: Optional parameters
    """
    with client.cursor() as cursor:
        cursor.execute(query, parameters)
        return [row[0] for row in cursor]


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

    Parameters:
        cursor: The psycopg Cursor to use
        query: The query to retrieve data
        parameters: Optional set of parameters passed to the cursor
    """
    with client.cursor() as cursor:
        cursor.execute(query, parameters)
        row = namedtuple('Row', [col.name for col in cursor.description])
        return [row(*r) for r in cursor]