Designing an URL Shortener

Question

Recently I was assigned a URL shortener design problem while interviewing with a well known organization. Here is the exact problem statement:

1. Build a simple URL shortener service that will accept a URL as an argument over a REST API and return a shortened URL as a result.

2. The URL and shortened URL should be stored in memory by applicaon.

   a. [BONUS] Instead of in memory, store these things in a text file.

3. If I again ask for the same URL, it should give me the same URL as it gave before instead of generating a new one.

4. [BONUS] Put this application in a Docker image by writing a Dockerfile and provide the docker image link along with the source code link.

The problem looked simpler and I did all bonus points as well but still didn't hear back from them (not even the feedback).

Here is my approach:

• As they have asked for multiple storage mechanisms, I created a factory for key-value stores and implemented memory and file based stores.
• I could think of two approaches for shortening the URLs: 1) Hashing 2) Base 62 of a counter. Base 62 seemed proper approach but they have added a requirement that for same long URL, the same shortened URL should be returned. As Base 62 works with an autoincrement counter, achieving this requirement without extra memory/storage is not possible. So I went with hashing the long URLs (I know there are chances of collisions, I mentioned this as a trade-off).

Here is my base key-value store:

kvstore_base.py

from abc import ABC, abstractmethod


class KVStoreBase(ABC):
    """
    Base key-value store class that provides abstract methods required for
    any storage mechanism used to store shortened URLs.
    """

    def __init__(self, store):
        """
        Constructor.

        :param store: Object of any storage mechanism.
        """
        self.store = store

    @abstractmethod
    def __getitem__(self, key: str) -> str:
        pass

    @abstractmethod
    def __setitem__(self, key: str, value: str) -> None:
        pass

    @abstractmethod
    def __contains__(self, key: str) -> bool:
        pass

file_store.py

import io

from .kvstore_base import KVStoreBase


class FileStore(KVStoreBase):
    """
    Memory store, which uses 'dict' data structure to
    store the key-value pairs.
    """

    def __init__(self):
        super().__init__(open("db.txt", "a+"))

    def __getitem__(self, key: str) -> str:
        """
        Get value for the given key from dict
        :param key: Key for which the value is to be retrieved.
        :return: value
        """
        # iterate over file and search for given key.
        try:
            # move pointer to initial position in file
            self.store.seek(0, io.SEEK_SET)
            for line in self.store:
                suffix, long_url = line.split()
                if suffix == key:
                    return long_url
        except Exception as err:
            print(str(err))
        return ""

    def __setitem__(self, key: str, value: str) -> None:
        """
        Set value for the given key into dict.
        :param key: Key to be added.
        :param value: Value corresponding to the key.
        :return: None
        """
        # move pointer to the end of the file for writing.
        self.store.seek(0, io.SEEK_END)
        if self.store.tell() != 0:
            self.store.write(f"\n{key} {value}")
        else:
            self.store.write(f"{key} {value}")
        self.store.flush()

    def __contains__(self, key: str) -> bool:
        """
        Check whether the given key is present in the dict.
        :param key: Key whose presence is to be checked.
        :return: True if key is present, False otherwise.
        """
        return True if self.__getitem__(key) else False

    def __del__(self) -> None:
        """
        Free file resource.
        """
        self.store.close()

kvstore_factory.py

from .memory_store import MemoryStore
from .file_store import FileStore
from .kvstore_base import KVStoreBase


class Factory:
    """
    KV Store factory to instantiate storage objects.
    """

    kvstore_map: dict = {"MEMORY": MemoryStore, "FILE": FileStore}

    @staticmethod
    def get_instance(instance_type: str) -> KVStoreBase:
        """
        Instantiate given KV store class dynamically.
        :param instance_type: Type of KV store to be instantiated.
        :return: KV store object.
        """
        try:
            return Factory.kvstore_map[instance_type]()
        except KeyError:
            raise Exception("Invalid instance requested.")

I used FastAPI to create endpoints, here is the two main routes (one for shortening the URL and one for redirecting):

routes/urls.py

import hashlib
from urllib.parse import urlparse

from fastapi import APIRouter, HTTPException, Request

from ..utils.store_connector import StoreConnector
from ..models.urls import Url
from starlette.responses import RedirectResponse

router = APIRouter()
store = StoreConnector().store


def _get_base_url(endpoint_url: str) -> str:
    """
    Extract base url from any endpoint URL.
    :param endpoint_url: Endpoint URL from which the base URL is to be extracted.
    :return: Base URL
    """
    return f"{urlparse(endpoint_url).scheme}://{urlparse(endpoint_url).hostname}:{urlparse(endpoint_url).port}"


@router.post("/shorten", tags=["URLs"])
async def shorten(url_obj: Url, request: Request) -> Url:
    """
    Shorten the given long URL.
    :param request: request object
    :param url_obj: URL object
    :return: shortened URL.
    """
    suffix = hashlib.sha256(url_obj.url.encode("utf-8")).hexdigest()[:8]
    if suffix not in store:
        # store short-url-suffix: long-url into data store.
        store[suffix] = url_obj.url

    return Url(url=f"{_get_base_url(request.url_for('shorten'))}/{suffix}")


@router.get("/{suffix}", tags=["URLs"])
async def redirect(suffix: str) -> RedirectResponse:
    """
    Redirect to long URL for the given URL ID.
    :param suffix: URL ID for the corresponding long URL.
    :return: Long URL.
    """
    long_url = store[suffix]
    if long_url:
        # return permanent redirect so that browsers store this in their cache.
        response = RedirectResponse(url=long_url, status_code=301)
        return response
    raise HTTPException(status_code=404, detail="Short URL not found.")

I exposed the storage mechanism to docker environment variable and created a storage connector class, which calls the factory based on the user configuration.

store_connector.py

import os
import sys
from sys import stderr

from .singleton import Singleton
from ..kvstores.kvstore_factory import Factory
from ..kvstores.kvstore_base import KVStoreBase


class StoreConnector(metaclass=Singleton):
    """
    Key Value store singleton class that can be used across all modules.
    """

    def __init__(self):
        try:
            store_type = os.environ.get("STORE_TYPE", "MEMORY")
            self._store = Factory.get_instance(store_type)
        except KeyError as ex:
            print(ex, file=stderr)
            # one of the required environment variable is not set
            print(
                "One of the required environment variables is not set",
                file=stderr,
            )
            sys.exit(1)
        except Exception as ex:
            print(ex, file=stderr)
            sys.exit(1)

    @property
    def store(self) -> KVStoreBase:
        return self._store

Here is my directory structure:

    url-shortener/
├── Dockerfile
├── LICENSE
├── README.md
├── requirements.txt
├── src
│   ├── __init__.py
│   ├── kvstores
│   │   ├── __init__.py
│   │   ├── file_store.py
│   │   ├── kvstore_base.py
│   │   ├── kvstore_factory.py
│   │   └── memory_store.py
│   ├── main.py
│   ├── models
│   │   ├── __init__.py
│   │   └── urls.py
│   ├── routes
│   │   ├── __init__.py
│   │   └── urls.py
│   └── utils
│       ├── singleton.py
│       └── store_connector.py
├── start.sh
└── tests
    ├── __init__.py
    ├── test_file_store.py
    ├── test_kvstore_factory.py
    ├── test_memory_store.py
    ├── test_routes.py
    └── test_store.py

In the instructions, they emphasized the following points:

• Readability of code
• Tests - Unit tests definitely and more if you can think of
• A good structure to your code and well written file & variable names etc.

So, which of those points my code lacks?

Answer 1

Your code looks nice and readable in my opinion; I can understand most of what's going on. A few points I'm unsure about, however:

(1) Your exception handling seems a bit weird to me

In filestore.py, you catch all exceptions in your __getitem__ method. This is generally bad practice — if there's a specific kind of exception that you're anticipating, then you should catch that specific exception only. There are some situations where you might want to catch and suppress all exceptions, but they're rare, and (to me at least) it's not evident why this would be one of them. There's nothing in the docstring or anything to explain why you're catching it and just printing it out. Rather than catching all errors in __getitem__ and returning an empty string if you weren't successful, I would catch KeyErrors in routes/urls.py

Similarly, in kvstore_factory.py, it's unclear to me why you'd catch a KeyError exception — a specific kind of exception, which is informative about what might have caused it — and replace it with a general Exception. If you want to add some useful additional information into the error message, I would do either this:

# kvstore_factory.py

class Factory:
    kvstore_map: dict = {"MEMORY": MemoryStore, "FILE": FileStore}

    @staticmethod
    def get_instance(instance_type: str) -> KVStoreBase:
        try:
            return Factory.kvstore_map[instance_type]()
        except KeyError as err:
            raise KeyError("Invalid instance requested.") from err

Or this:

# kvstore_factory.py

class InvalidInstanceRequest(KeyError):
    """
    Raised when a user requests an instance
    that does not exist in Factory.kvstore_map
    """


class Factory:
    kvstore_map: dict = {"MEMORY": MemoryStore, "FILE": FileStore}

    @staticmethod
    def get_instance(instance_type: str) -> KVStoreBase:
        try:
            return Factory.kvstore_map[instance_type]()
        except KeyError as err:
            raise InvalidInstanceRequest("Invalid instance requested.") from err

Lastly, I'm confused about why, in store_connector.py, you're catching exceptions and then calling sys.exit. Why not just re-raise the original exception, like so? I think it would lead to a more readable traceback. I would take out the second except clause altogether — if all you're doing is printing it and then calling sys.exit(), you may as well not catch it at all and let the exception end the programme.

# store_connector.py

import os

from .singleton import Singleton
from ..kvstores.kvstore_factory import Factory


class StoreConnector(metaclass=Singleton):
    def __init__(self):
        try:
            store_type = os.environ.get("STORE_TYPE", "MEMORY")
            self._store = Factory.get_instance(store_type)
        except KeyError as ex:
           raise KeyError("One of the required environment variables is not set") from ex

(2) Two minor points

Firstly, in kvstore_factory.py, you could have much richer type hints. I'm also not quite sure why you're using a staticmethod, when a classmethod would seem more elegant here.

# kvstore_factory.py

from typing import ClassVar
from .memory_store import MemoryStore
from .file_store import FileStore
from .kvstore_base import KVStoreBase


class Factory:
    kvstore_map: ClassVar[dict[str, KVStoreBase]] = {"MEMORY": MemoryStore, "FILE": FileStore}

    @classmethod
    def get_instance(cls, instance_type: str) -> KVStoreBase:
        try:
            return cls.kvstore_map[instance_type]()
        except KeyError as ex:
            # exception handling as discussed above 

# -- snip --

Secondly, you can simplify your __contains__ method in file_store.py by simply writing return bool(self[key])

(3) I would have put some of the logic in routes/urls.py elsewhere

The exact manner in which you're hashing the urls seems like an implementation detail. It seems strange to me that this logic is in routes/urls.py. I might rewrite your code like this:

kvstore_base.py:

# kvstore_base.py

import hashlib
from abc import ABC, abstractmethod
from collections.abc import MutableMapping
from functools import cache

@MutableMapping.register
class KVStoreBase(ABC):
    def __init__(self, store):
        self.store = store

    def get_shortened_url(self, long_url: str) -> str:
        shortened_url = self.shorten_url(long_url)
        self[shortened_url] = long_url
        return shortened_url

    @staticmethod
    @cache
    def shorten_url(url: str) -> str:
        return hashlib.sha256(url.encode("utf-8")).hexdigest()[:8]

    @abstractmethod
    def __getitem__(self, key: str) -> str: ...

    @abstractmethod
    def __setitem__(self, key: str, value: str) -> None: ...

    @abstractmethod
    def __contains__(self, key: str) -> bool: ...

routes/urls.py:

# routes/urls.py

from urllib.parse import urlparse

from fastapi import APIRouter, HTTPException, Request

from ..utils.store_connector import StoreConnector
from ..models.urls import Url
from starlette.responses import RedirectResponse

router = APIRouter()
store = StoreConnector().store


def _get_base_url(endpoint_url: str) -> str:
    return f"{urlparse(endpoint_url).scheme}://{urlparse(endpoint_url).hostname}:{urlparse(endpoint_url).port}"


@router.post("/shorten", tags=["URLs"])
async def shorten(url_obj: Url, request: Request) -> Url:
    suffix = store.get_shortened_url(url_obj.url)
    return Url(url=f"{_get_base_url(request.url_for('shorten'))}/{suffix}")

# -- snip --

Alternative refactoring of KVStore class, if you don't believe hashing details belong in the base class KVStore class:

# kvstore_base.py

import hashlib
from abc import ABC, abstractmethod
from collections.abc import MutableMapping, Callable
from functools import cache
from typing import TypeVar

@MutableMapping.register
class KVStoreBase(ABC):
    def __init__(self, store):
        self.store = store

    @abstractmethod
    def __getitem__(self, key: str) -> str: ...

    @abstractmethod
    def __setitem__(self, key: str, value: str) -> None: ...

    @abstractmethod
    def __contains__(self, key: str) -> bool: ...


T = TypeVar('T', bound='URLStoreMixin')


class URLStoreMixin:
    """
    Mixin class for extending the capabilities of classes inheriting from KVStoreBase,
    useful for situations in which the KVStore is being used
    to store mappings of shortened URLs to long URLs
    """

    # type hint so mypy doesn't get confused by this being a Mixin class
    __setitem__: Callable[[T, str, str], None]

    def get_shortened_url(self, long_url: str) -> str:
        shortened_url = self.shorten_url(long_url)
        self[shortened_url] = long_url
        return shortened_url

    @staticmethod
    @cache
    def shorten_url(url: str) -> str:
        return hashlib.sha256(url.encode("utf-8")).hexdigest()[:8]

Then have your MemoryStore and FileStore objects solely inherit from KVStoreBase, as you were before, but create MemoryURLStore and FileURLStore classes like so:

class MemoryURLStore(MemoryStore, URLStoreMixin):
    pass

class FileURLStore(FileStore, URLStoreMixin):
    pass