152 lines
3.8 KiB
Python
152 lines
3.8 KiB
Python
"""
|
|
Main code
|
|
"""
|
|
|
|
import random
|
|
import string
|
|
from dataclasses import dataclass
|
|
from datetime import datetime
|
|
from typing import Optional
|
|
|
|
import sqlalchemy # type: ignore[import]
|
|
|
|
from slinky import db
|
|
|
|
|
|
@dataclass
|
|
class Shortcode:
|
|
"""
|
|
Simple dataclass to allow for typing of Shortcodes
|
|
"""
|
|
|
|
id: int # pylint: disable=invalid-name
|
|
shortcode: str
|
|
url: str
|
|
fixed_views: int
|
|
expiry: str
|
|
|
|
|
|
def random_string(length: int = 4) -> str:
|
|
"""
|
|
Create a random, alphanumeric string of the length specified
|
|
|
|
Args:
|
|
length (int, optional): length of string to generate. Defaults to 4.
|
|
|
|
Returns:
|
|
str: random alphanumeric string
|
|
"""
|
|
allowed_chars: str = string.ascii_letters + string.digits
|
|
|
|
return ''.join(random.SystemRandom().choice(allowed_chars) for _ in range(length))
|
|
|
|
|
|
class Slinky:
|
|
"""
|
|
Class for Slinky
|
|
"""
|
|
|
|
def __init__(self, db_url: str) -> None:
|
|
self.db = db.ShortcodeDB(db_url) # pylint: disable=invalid-name
|
|
self.session = self.db.session()
|
|
|
|
def add(
|
|
self,
|
|
url: str,
|
|
length: int = 4,
|
|
fixed_views: int = -1,
|
|
expiry: datetime = datetime.max,
|
|
) -> str:
|
|
"""
|
|
Add a shortcode to the DB
|
|
|
|
Args:
|
|
url (str): URL to redirect to
|
|
fixed_views (int, optional): number of views to serve before expiring.
|
|
Defaults to 0 (no limit).
|
|
expiry (int, optional): date of expiry. Defaults to 0 (no limit).
|
|
|
|
Returns:
|
|
str: shortcode for the redirect
|
|
"""
|
|
shortcode = random_string(length=length)
|
|
|
|
if self.get_by_shortcode(shortcode).url:
|
|
raise ValueError(f'Shortcode {shortcode} already exists')
|
|
|
|
dbentry = db.ShortURL(
|
|
shortcode=shortcode,
|
|
url=url,
|
|
fixed_views=fixed_views,
|
|
expiry=expiry,
|
|
)
|
|
self.session.add(dbentry)
|
|
self.session.commit()
|
|
self.session.close()
|
|
|
|
return shortcode
|
|
|
|
def get_by_shortcode(self, shortcode: str) -> Shortcode:
|
|
"""
|
|
Return a Shortcode object for a given shortcode
|
|
|
|
Args:
|
|
shortcode (str): the shortcode to look up
|
|
|
|
Returns:
|
|
Shortcode: full Shortcode object for the given shortcode
|
|
"""
|
|
entry = self.session.query(db.ShortURL).filter_by(shortcode=shortcode).first()
|
|
|
|
if entry:
|
|
ret_sc = Shortcode(
|
|
entry.id,
|
|
entry.shortcode,
|
|
entry.url,
|
|
entry.fixed_views,
|
|
entry.expiry,
|
|
)
|
|
self.session.close()
|
|
return ret_sc
|
|
return Shortcode(0, '', '', 0, '1970-01-01 00:00:00.000000')
|
|
|
|
def remove_view(self, sc_id: int) -> None:
|
|
"""
|
|
Reduce the fixed views count by one
|
|
|
|
Args:
|
|
id (int): ID of the DB entry to reduce the fixed_views count
|
|
"""
|
|
self.session.query(db.ShortURL).filter_by(id=sc_id).update(
|
|
{db.ShortURL.fixed_views: db.ShortURL.fixed_views - 1}
|
|
)
|
|
self.session.commit()
|
|
self.session.close()
|
|
|
|
def get_all(self) -> list[Shortcode]:
|
|
"""
|
|
Return a Shortcode object for a given shortcode
|
|
|
|
Args:
|
|
shortcode (str): the shortcode to look up
|
|
|
|
Returns:
|
|
Shortcode: full Shortcode object for the given shortcode
|
|
"""
|
|
shortcodes = list(self.session.query(db.ShortURL).all())
|
|
self.session.close()
|
|
return shortcodes
|
|
|
|
def delete_by_shortcode(self, shortcode: str) -> None:
|
|
"""
|
|
Delete shortcode entry
|
|
|
|
Args:
|
|
shortcode (str): Shortcode of entry to delete
|
|
"""
|
|
entry = self.session.query(db.ShortURL).filter_by(shortcode=shortcode).first()
|
|
|
|
self.session.delete(entry)
|
|
self.session.commit()
|
|
self.session.close()
|