Coverage for src / links / crud.py: 100%

1from sqlalchemy.ext.asyncio import AsyncSession 

2from sqlalchemy.future import select 

3from . import models 

4from datetime import datetime, timezone 

5import secrets 

6import string 

7from typing import Optional 

8from .exceptions import NotUniqueAliasError 

9 

10 

11async def create_link( 

12 db: AsyncSession, 

13 original_url: str, 

14 custom_alias: Optional[str] = None, 

15 expire_at: Optional[datetime] = None, 

16 user_id: Optional[int] = None 

17) -> models.Link: 

18 if custom_alias: 

19 existing = await db.execute( 

20 select(models.Link).filter(models.Link.custom_alias == custom_alias) 

21 ) 

22 if existing.scalars().first(): 

23 raise NotUniqueAliasError(custom_alias) 

24 

25 short_id = custom_alias or generate_short_id() 

26 link = models.Link( 

27 original_url=original_url, 

28 short_id=short_id, 

29 user_id=user_id, 

30 expire_at=expire_at 

31 ) 

32 db.add(link) 

33 await db.commit() 

34 await db.refresh(link) 

35 return link 

36 

37 

38async def get_link_by_short_id(db: AsyncSession, short_id: str) -> Optional[models.Link]: 

39 result = await db.execute( 

40 select(models.Link).filter(models.Link.short_id == short_id) 

41 ) 

42 return result.scalars().first() 

43 

44 

45async def increment_click_count(db: AsyncSession, link: models.Link) -> None: 

46 link.click_count += 1 

47 await db.commit() 

48 

49 

50async def search_links( 

51 db: AsyncSession, 

52 original_url: str, 

53 user_id: int 

54) -> list[models.Link]: 

55 result = await db.execute( 

56 select(models.Link) 

57 .filter(models.Link.original_url.contains(original_url)) 

58 .filter(models.Link.user_id == user_id) 

59 ) 

60 return result.scalars().all() 

61 

62 

63async def update_link( 

64 db: AsyncSession, 

65 link: models.Link, 

66 new_url: str 

67) -> models.Link: 

68 link.original_url = new_url 

69 await db.commit() 

70 await db.refresh(link) 

71 return link 

72 

73 

74async def delete_link(db: AsyncSession, link: models.Link) -> None: 

75 await db.delete(link) 

76 await db.commit() 

77 

78 

79def generate_short_id(length=6) -> str: 

80 return ''.join(secrets.choice(string.ascii_letters + string.digits) for _ in range(length))