PyConES24/src/retailtwin/cli/db.py
2024-05-08 21:55:08 +00:00

110 lines
2.7 KiB
Python

"""
Several utilities to fetch data from the database
"""
import polars as pl
from rich.table import Table
from typing import Tuple
def fetch_from_db(db_uri: str, location: int) -> Tuple[pl.DataFrame]:
"""Fetcn basic information about inventory from the database
Args:
db_uri (str): SQLAlchemy db URI
location (int): Location ID for the store or warehouse
Returns:
Tuple[pl.DataFrame]: Two polars dataframe continig all items,
and the local stock respectively
"""
items = (
query_items(db_uri)
.select(pl.all())
.select([pl.col("upc"), pl.col("name"), pl.col("package")])
.with_columns(pl.col("upc").cast(str))
)
local_stock = query_local_batches(db_uri, location).with_columns(
pl.col("upc").cast(str),
pl.col("batch").cast(str),
pl.col("quantity").cast(str),
pl.col("received").cast(str),
pl.col("best_until").cast(str),
)
return items, local_stock
def df_to_table(df: pl.DataFrame, title="Items") -> Table:
"""Process a Polars dataframe and create a Rich table from it
Args:
df (pl.DataFrame): The polars dataframe
title (str, optional): Title for the displayed table. Defaults to "Items".
Returns:
Table: The Rich table to be displayed
"""
table = Table(title=title)
for column in df.columns:
table.add_column(str(column))
for row in df.iter_rows():
table.add_row(*row)
return table
def query_items(db_uri: str) -> pl.DataFrame:
"""Query all available items from the database that can be kept in stock
Args:
db_uri (str): SQLAlchemy database URI
Returns:
pl.DataFrame: Polars dataframe with the items
"""
return pl.read_database(
"select * from items where current = true", db_uri, engine="adbc"
)
def query_local_batches(db_uri: str, location: int) -> pl.DataFrame:
"""Query all batches stored in location
Args:
db_uri (str): SQLAlchemy database URI
location (int): Id of the current location
Returns:
pl.DataFrame: Polars Dataframe with the batches in stock
"""
return pl.read_database(
f"select * from inventory where location = {location}",
db_uri,
engine="adbc",
)
def query_warehouse_stock(db_uri: str, item: int) -> pl.DataFrame:
"""Returns stock available for an item in all warehouses
Args:
db_uri (str): SQLAlchemy database URI
item (int): Id of the item
Returns:
pl.DataFrame: Polars Dataframe with the stock stock
"""
return pl.read_database(
f"select * from warehouse_stock where upc = {item}",
db_uri,
engine="adbc",
)
def sync_pos_schema():
pass