shortener.py

from __future__ import print_function
import flask
import os
from hashlib import md5
from base64 import b64encode
from pynamodb.models import Model
from pynamodb.indexes import GlobalSecondaryIndex, AllProjection
from pynamodb.attributes import UnicodeAttribute


class Url(Model):
    class Meta:
        table_name = "shortened-urls"
        host = "http://127.0.0.1:8000"
        region = "us-west-2"

    class CodeIndex(GlobalSecondaryIndex):
        class Meta:
            read_capacity_units = 1
            write_capacity_units = 1
            projection = AllProjection()
        code = UnicodeAttribute(hash_key=True)

    url = UnicodeAttribute(hash_key=True)
    code = UnicodeAttribute()
    code_index = CodeIndex()

    def save(self, **kwargs):
        """
        Generates the shortened code before saving
        """
        self.code = b64encode(
            md5(self.url.encode('utf-8')).hexdigest()[-4:].encode('utf-8')
        ).decode('utf-8').replace('=', '').replace('/', '_')
        super(Url, self).save(**kwargs)

app = flask.Flask(__name__)
app.config.update(DEBUG=True)


@app.route('/')
def index():
    return flask.render_template(
        "index.html",
        url_root = flask.request.url_root)


@app.route('/shorten/<path:url>')
def shorten(url):
    model = Url(url)
    model.save()
    return flask.Response(model.code)


@app.route('/<path:code>')
def resolve(code):
    # next() in Python3 is __next__()
    try:
        model = Url.code_index.query(code).next()
    except AttributeError:
        model = Url.code_index.query(code).__next__()
    except StopIteration:
        flask.abort(404)
    finally:
        return flask.redirect(model.url)

if __name__ == "__main__":
    if not Url.exists():
        print("Creating table...")
        Url.create_table(wait=True, read_capacity_units=1, write_capacity_units=1)
    app.run(host='0.0.0.0')