package main import ( "database/sql" "fmt" "log" "math/rand" "net/http" "os" "time" "github.com/gorilla/mux" _ "github.com/lib/pq" ) type URLShortener struct { DB *sql.DB } type URL struct { ID int OriginalURL string ShortCode string CreatedAt time.Time Visits int } func main() { dbHost := getEnv("DB_HOST", "localhost") dbPort := getEnv("DB_PORT", "5432") dbUser := getEnv("DB_USER", "postgres") dbPassword := getEnv("DB_PASSWORD", "postgres") dbName := getEnv("DB_NAME", "urlshortener") connStr := fmt.Sprintf("host=%s port=%s user=%s password=%s dbname=%s sslmode=disable", dbHost, dbPort, dbUser, dbPassword, dbName) db, err := sql.Open("postgres", connStr) if err != nil { log.Fatalf("Failed to connect to the database: %v", err) } defer db.Close() err = db.Ping() if err != nil { log.Fatalf("Failed to ping the database: %v", err) } err = initializeDB(db) if err != nil { log.Fatalf("Failed to initialize the DB: %v", err) } shortener := &URLShortener{DB: db} router := mux.NewRouter() router.PathPrefix("/static/").Handler(http.StripPrefix("/static/", http.FileServer(http.Dir("./static")))) router.HandleFunc("/", shortener.homeHandler).Methods("GET") router.HandleFunc("/shorten", shortener.shortenHandler).Methods("POST") router.HandleFunc("/stats", shortener.statsHandler).Methods("GET") router.HandleFunc("/{shortCode}", shortener.redirectHandler).Methods("GET") port := getEnv("PORT", "8080") log.Printf("Server starting on port http://localhost:%s", port) if err := http.ListenAndServe(":"+port, router); err != nil { log.Fatalf("Server failed to start: %v", err) } } func getEnv(key, fallback string) string { if value, exists := os.LookupEnv(key); exists { return value } return fallback } func initializeDB(db *sql.DB) error { _, err := db.Exec(` CREATE TABLE IF NOT EXISTS urls( id SERIAL PRIMARY KEY, original_url TEXT NOT NULL, short_code VARCHAR(10) UNIQUE NOT NULL, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, visits INTEGER DEFAULT 0 ) `) return err } func (s *URLShortener) homeHandler(w http.ResponseWriter, r *http.Request) { http.ServeFile(w, r, "./templates/index.html") } func (s *URLShortener) shortenHandler(w http.ResponseWriter, r *http.Request) { err := r.ParseForm() if err != nil { http.Error(w, "Error parsing form", http.StatusBadRequest) return } originalURL := r.FormValue("url") if originalURL == "" { http.Error(w, "URL is required", http.StatusBadRequest) return } shortcode := generateShortCode(6) _, err = s.DB.Exec( "INSERT INTO urls (original_url, short_code) VALUES ($1, $2)", originalURL, shortcode, ) if err != nil { log.Printf("Error inserting URL: %v", err) http.Error(w, "Error creating short url", http.StatusInternalServerError) return } if r.Header.Get("HX-Request") == "true" { baseURL := fmt.Sprintf("%s://%s", getScheme(r), r.Host) shortURL := fmt.Sprintf("%s/%s", baseURL, shortcode) w.Header().Set("Content-Type", "text/html") w.Write([]byte(fmt.Sprintf(`
Your shortened URL:
%s| Original URL | Short URL | Created | Visists |
|---|---|---|---|
| %s | %s | %s | %d |