from sqlalchemy.exc import SQLAlchemyError

from load.projects.cloud.loadtesting.db.tables import AmmoTable, StorageTable
from load.projects.cloud.loadtesting.logan import lookup_logger


class AmmoQueries:
    def __init__(self, session):
        self._session = session

    def get(self, ammo_id) -> AmmoTable:
        return self._session.query(AmmoTable).filter(AmmoTable.id == ammo_id).first()

    def get_by_name(self, name, bucket, folder_id):
        query = self._session.query(AmmoTable).join(StorageTable).filter(
            AmmoTable.folder_id == folder_id, StorageTable.id == AmmoTable.storage_id, StorageTable.bucket == bucket, AmmoTable.s3_name == name)
        return query.first()

    def add(self, ammo):
        try:
            self._session.add(ammo)
            lookup_logger().debug(f'Added ammo {ammo.id}')
        except SQLAlchemyError:
            lookup_logger().exception(f'Exception while adding ammo {ammo.id}.')
            raise
