SatisfactoryPlusCalculator/factorygame/data/fetch.py

#!/usr/bin/env python3

import click
from sqlalchemy import create_engine, select
from sqlalchemy.orm import Session

from factorygame.data.common import resource_needs_update, chose_resource
from .models import Base, ResourceFlow, Recipe
from .sfp import SatisfactoryPlus
from ..helper import click_prompt


@click.command()
@click.option("--result", is_flag=True, default=True)
@click.option("--debug", is_flag=True)
@click.option("--refetch", is_flag=True)
@click.argument("search")
def main(result: bool, debug: bool, refetch: bool, search: str):
    engine = create_engine("sqlite:///file.db", echo=debug)
    Base.metadata.create_all(bind=engine)
    if result and search:
        with Session(engine) as session:
            resource = chose_resource(session=session, resource_label=search, prompt=click_prompt).result()
            exists_in_db = resource is not None

            with SatisfactoryPlus(debug=debug) as data_provider:
                if resource is None:
                    ret = data_provider.search_for_resource(session=session, search=search)
                    if ret is None:
                        return
                    else:
                        resource, exists_in_db = ret

                refetch |= resource_needs_update(resource)
                if refetch:
                    if exists_in_db:
                        print("Deleting recipes for", resource.label)
                        with session.begin_nested():
                            for flow in session.scalars(
                                select(ResourceFlow).where(ResourceFlow.resource_id == resource.id)
                            ):
                                if flow.result_of:
                                    for flow2 in flow.result_of.ingredients:
                                        session.delete(flow2)
                                    for flow2 in flow.result_of.results:
                                        session.delete(flow2)
                                    session.delete(flow.result_of)
                        print("Refetching recipes for", resource.label)
                    else:
                        print("Fetching recipes for new resource", resource.label)

                    with session.begin_nested():
                        resource = data_provider.update_resource_recipes(session=session, resource=resource)

                session.refresh(resource)
                assert resource, "Resource must be set at this point"

                stmt = select(Recipe).join(Recipe.results).filter(ResourceFlow.resource_id == resource.id)
                for recipe in session.scalars(stmt):
                    print("Recipe:", recipe.describe())


if __name__ == "__main__":
    main()