import asyncio import json import logging import secrets import uuid from datetime import UTC, date, datetime from itertools import zip_longest from time import time from typing import Any from aiokafka.errors import KafkaError from url_shortener.adapters.aws.s3 import S3Client from url_shortener.adapters.kafka.kafka_client import KafkaClient from url_shortener.dtos import ShortUrl, ShortUrls from url_shortener.exceptions import ( NotAllowedDomainError, PathAlreadyInUseError, PathAndPathPrefixProvidedError, ) from url_shortener.repositories import PathsRepository from url_shortener.utils import chunks logger = logging.getLogger(__name__) class ShortUrlRedirectService: def __init__( self, repository: PathsRepository, kafka_client: KafkaClient, s3_client: S3Client, failed_events_s3_bucket: str, ) -> None: self.repository = repository self.kafka_client = kafka_client self.s3_client = s3_client self.failed_events_s3_bucket = failed_events_s3_bucket async def get_redirect_url(self, path: str) -> dict[str, Any] | None: return await self.repository.get_url(path) async def collect_analytics( self, path: str, redirect_url: str | None, is_personalized: bool, created_at: float, additional_attributes: dict[str, Any] | None, **kwargs: Any, ) -> None: analytics = { "path": path, "redirect_url": redirect_url, "is_personalized": is_personalized, "timestamp": int(time()), "additional_attributes": additional_attributes, "created_at": created_at, **kwargs, } try: await self.kafka_client.put(analytics) except KafkaError: self.s3_client.put_object( bucket=self.failed_events_s3_bucket, key=f"{date.today().isoformat()}/{uuid.uuid4()}.json", body=json.dumps(analytics), ) logger.warning("Failed to put analytics to Kafka, saved to S3") class ShortenUrlService: def __init__( self, repository: PathsRepository, allowed_domains: list[str], short_path_allowed_chars: str, ) -> None: self.repository = repository self.allowed_domains = allowed_domains self.short_path_allowed_chars = short_path_allowed_chars async def get_short_url( self, url: str, domain: str, path_length: int, path: str | None = None, path_prefix: str | None = None, additional_attributes: dict[str, Any] | None = None, ) -> ShortUrl: if domain not in self.allowed_domains: raise NotAllowedDomainError if path and path_prefix: raise PathAndPathPrefixProvidedError if path and not await self.repository.is_path_available(path): raise PathAlreadyInUseError is_personalized = not path path = path or await self._generate_path(path_prefix, path_length) created_at = int(time()) await self.repository.put_url( url, path, is_personalized, created_at=created_at, additional_attributes=additional_attributes, ) return ShortUrl( short_url=f"{domain}/{path}", is_personalized=is_personalized, url=url, created_at=datetime.fromtimestamp(created_at, UTC), additional_attributes=additional_attributes, ) async def get_short_urls( self, url: str, domain: str, path_length: int, path_prefix: str | None = None, urls_count: int = 1, additional_attributes: list[dict[str, Any]] | None = None, ) -> ShortUrls: if domain not in self.allowed_domains: raise NotAllowedDomainError paths = await self._generate_paths( path_length, path_prefix, urls_count=urls_count, ) created_at = int(time()) # 25 is fixed limit for write into dynamodb paths_chunks = chunks(paths, chunk_size=25) if additional_attributes: additional_attributes_chunks = chunks(additional_attributes, chunk_size=25) else: additional_attributes_chunks = [] tasks = [] for path_chunk, additional_attributes_chunk in zip_longest( paths_chunks, additional_attributes_chunks ): tasks.append( self.repository.put_urls( url, created_at, path_chunk, additional_attributes=additional_attributes_chunk, # type: ignore ) ) await asyncio.gather(*tasks) short_urls = [ { "short_url": f"{domain}/{path}", "additional_attributes": path_additional_attributes, } for path, path_additional_attributes in zip_longest( paths, additional_attributes or [] ) ] return ShortUrls.model_validate( {"short_urls": short_urls, "url": url, "created_at": created_at} ) async def _generate_paths( self, path_length: int, path_prefix: str | None, urls_count: int = 1, ) -> list[str]: paths = set() # Possible improvement here, split paths check using asyncio.gather while len(paths) < urls_count: # Generate paths until we will have needed amount of free paths path_strings = [ self._generate_path_string(path_length, path_prefix) for _ in range( min(urls_count - len(paths), 100) ) # 100 is a fixed limit for read from dynamodb ] urls = await self.repository.get_urls(path_strings) existing_paths = {url["path"] for url in urls} for path in path_strings: # Use only available paths if path not in existing_paths: paths.add(path) return list(paths) async def _generate_path(self, path_prefix: str | None, path_length: int) -> str: while True: path_string = self._generate_path_string( path_length=path_length, path_prefix=path_prefix, ) if await self.repository.is_path_available(path_string): return path_string def _generate_path_string(self, path_length: int, path_prefix: str | None) -> str: path_prefix = f"{path_prefix}-" if path_prefix else "" path = "".join( secrets.choice(self.short_path_allowed_chars) for _ in range(path_length) ) return path_prefix + path class DeletePathsService: def __init__(self, repository: PathsRepository): self.repository = repository async def delete_paths(self, paths: list[str]) -> None: paths_chunks = chunks(paths, chunk_size=25) await asyncio.gather(*map(self.repository.delete_paths, paths_chunks)) class UpdatePathService: def __init__(self, repository: PathsRepository, allowed_domains: list[str]) -> None: self.repository = repository self.allowed_domains = allowed_domains async def update_path(self, path: str, update_attributes: dict[str, Any]) -> None: domain = update_attributes.get("domain") if domain and domain not in self.allowed_domains: raise NotAllowedDomainError await self.repository.update_path(path, update_attributes)