Source code for allensdk.internal.api.queries.behavior_lims_queries

from typing import Optional, List
import pandas as pd

from allensdk.internal.api import PostgresQueryMixin
import logging
from allensdk.internal.api.queries.utils import (
    build_in_list_selector_query)


def foraging_id_map_from_behavior_session_id(
        lims_engine: PostgresQueryMixin,
        behavior_session_ids: List[int],
        logger: Optional[logging.RootLogger] = None) -> pd.DataFrame:
    """
    Returns DataFrame with two columns:
        foraging_id
        behavior_session_id

    Parameters
    ----------
    lims_engine: PostgresQueryMixin
        Means of connecting to the LIMS database

    behavior_session_ids: List[int]
        List of behavior_session_ids for which we want
        the foraging_id
    """

    behav_ids = build_in_list_selector_query("id",
                                             behavior_session_ids,
                                             operator="AND")
    forag_ids_query = f"""
            SELECT foraging_id, id as behavior_session_id
            FROM behavior_sessions
            WHERE foraging_id IS NOT NULL
            {behav_ids};
            """
    if logger is not None:
        logger.debug("get_foraging_ids_from_behavior_session query: \n"
                     f"{forag_ids_query}")
    foraging_id_map = lims_engine.select(forag_ids_query)

    if logger is not None:
        logger.debug(f"Retrieved {len(foraging_id_map)} foraging ids for"
                     " behavior stage query. "
                     f"Ids = {foraging_id_map.foraging_id}")
    return foraging_id_map


def stimulus_pickle_paths_from_behavior_session_ids(
        lims_connection: PostgresQueryMixin,
        behavior_session_id_list: List[int]) -> pd.DataFrame:
    """
    Get a DataFrame mapping behavior_session_id to
    stimulus_pickle_path

    Parameters
    ----------
    lims_connection: PostgresQueryMixin

    behavior_session_id_list: List[int]

    Returns
    -------
    beh_to_path: pd.DataFrame
        with columns
            behavior_session_id
            pkl_path
    """

    query = f"""
    SELECT
      beh.id as behavior_session_id
      ,wkf.storage_directory || wkf.filename as pkl_path
    FROM behavior_sessions AS beh
    JOIN well_known_files AS wkf
      ON wkf.attachable_id = beh.id
    JOIN well_known_file_types as wkft
      ON
      wkf.well_known_file_type_id = wkft.id
    WHERE
        wkft.name = 'StimulusPickle'
      AND
        wkf.attachable_type = 'BehaviorSession'
        {build_in_list_selector_query(
        operator='AND',
        col='beh.id',
        valid_list=behavior_session_id_list)}
    """

    beh_to_path = lims_connection.select(query)
    return beh_to_path