Skip to content

Instantly share code, notes, and snippets.

@decagondev
Created July 18, 2024 00:59
Show Gist options
  • Save decagondev/2b0e8ac26b9aebf0845919d98f9aef9e to your computer and use it in GitHub Desktop.
Save decagondev/2b0e8ac26b9aebf0845919d98f9aef9e to your computer and use it in GitHub Desktop.
# app.py
from flask import Flask, request, jsonify, redirect, g
import sqlite3
import string
import random
import os

app = Flask(__name__)
DATABASE = 'urls.db'

def get_db():
    db = getattr(g, '_database', None)
    if db is None:
        db = g._database = sqlite3.connect(DATABASE)
    return db

@app.teardown_appcontext
def close_connection(exception):
    db = getattr(g, '_database', None)
    if db is not None:
        db.close()

def init_db():
    with app.app_context():
        db = get_db()
        with app.open_resource('schema.sql', mode='r') as f:
            db.cursor().executescript(f.read())
        db.commit()

def generate_shortcode(length=6):
    characters = string.ascii_letters + string.digits
    return ''.join(random.choice(characters) for _ in range(length))

@app.route('/generate_shortcode', methods=['POST'])
def create_shortcode():
    try:
        data = request.get_json()
        long_url = data['url']
        shortcode = generate_shortcode()
        
        db = get_db()
        db.execute('INSERT INTO url_map (shortcode, url) VALUES (?, ?)', (shortcode, long_url))
        db.commit()
        
        return jsonify({'shortcode': shortcode}), 201
    except Exception as e:
        return jsonify({'error': str(e)}), 400

@app.route('/<shortcode>', methods=['GET'])
def get_url(shortcode):
    try:
        db = get_db()
        cur = db.execute('SELECT url FROM url_map WHERE shortcode = ?', (shortcode,))
        row = cur.fetchone()
        if row:
            return jsonify({'URL': row[0]}), 200
        else:
            return jsonify({'error': 'Shortcode not found'}), 404
    except Exception as e:
        return jsonify({'error': str(e)}), 400

if __name__ == '__main__':
    if not os.path.exists(DATABASE):
        init_db()
    app.run(debug=True)
-- schema.sql
DROP TABLE IF EXISTS url_map;

CREATE TABLE url_map (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    shortcode TEXT NOT NULL UNIQUE,
    url TEXT NOT NULL
);
<!-- index.html -->
<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>URL Shortener</title>
    <style>
        body { font-family: Arial, sans-serif; margin: 20px; }
        input, button { padding: 10px; margin: 5px; }
    </style>
</head>
<body>
    <h1>URL Shortener</h1>
    <input type="text" id="url" placeholder="Enter URL">
    <button onclick="generateShortcode()">Shorten</button>
    <div id="shortcode"></div>
    <input type="text" id="shortcode_input" placeholder="Enter Shortcode">
    <button onclick="getUrl()">Get URL</button>
    <div id="original_url"></div>

    <script>
        async function generateShortcode() {
            const url = document.getElementById('url').value;
            const response = await fetch('/generate_shortcode', {
                method: 'POST',
                headers: {
                    'Content-Type': 'application/json',
                },
                body: JSON.stringify({ url }),
            });
            const data = await response.json();
            document.getElementById('shortcode').innerText = data.shortcode ? `Shortcode: ${data.shortcode}` : `Error: ${data.error}`;
        }

        async function getUrl() {
            const shortcode = document.getElementById('shortcode_input').value;
            const response = await fetch(`/${shortcode}`);
            const data = await response.json();
            document.getElementById('original_url').innerText = data.URL ? `URL: ${data.URL}` : `Error: ${data.error}`;
        }
    </script>
</body>
</html>
# test_app.py
import unittest
import json
from app import app, init_db, get_db

class URLShortenerTestCase(unittest.TestCase):
    def setUp(self):
        self.app = app.test_client()
        with app.app_context():
            init_db()
    
    def tearDown(self):
        with app.app_context():
            db = get_db()
            db.execute('DELETE FROM url_map')
            db.commit()

    def test_generate_shortcode(self):
        response = self.app.post('/generate_shortcode', data=json.dumps({'url': 'http://example.com'}), content_type='application/json')
        self.assertEqual(response.status_code, 201)
        data = json.loads(response.data)
        self.assertIn('shortcode', data)

    def test_get_url(self):
        with app.app_context():
            db = get_db()
            db.execute('INSERT INTO url_map (shortcode, url) VALUES (?, ?)', ('abc123', 'http://example.com'))
            db.commit()
        
        response = self.app.get('/abc123')
        self.assertEqual(response.status_code, 200)
        data = json.loads(response.data)
        self.assertEqual(data['URL'], 'http://example.com')

    def test_get_url_not_found(self):
        response = self.app.get('/notexist')
        self.assertEqual(response.status_code, 404)
        data = json.loads(response.data)
        self.assertIn('error', data)

if __name__ == '__main__':
    unittest.main()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment