from sqlalchemy import select, insert from sqlalchemy.exc import IntegrityError from app.utils import encoder from app.models import Url from app.exceptions import ShortcodeConflict, ShortcodeNotFound class UrlService: def __init__(self, db): self.db = db async def get_url_by_shortcode(self, shortcode: str) -> str: query = select(Url.url).where(Url.shortcode == shortcode) result = await self.db.fetch_val(query) if result is None: raise ShortcodeNotFound(f"Shortcode '{shortcode}' not found") return result async def generate_shortcode_and_save(self, url: str): try: # Get the next ID from the sequence next_id = await self.db.fetch_val("SELECT nextval('urls_id_seq')") shortcode = encoder.encode(next_id) # Insert the new URL entry insert_stmt = insert(Url).values(id=next_id, url=url, shortcode=shortcode) await self.db.execute(insert_stmt) return {"shortcode": shortcode, "url": url} except IntegrityError as e: if 'shortcode' in str(e): raise ShortcodeConflict(f"Shortcode '{shortcode}' already in use") from None raise