diff --git a/main.go b/main.go index 8fc4dbd..bc86260 100644 --- a/main.go +++ b/main.go @@ -98,21 +98,27 @@ func ShortenHandler(w http.ResponseWriter, r *http.Request) { return } - var slug string - err := db.QueryRow("SELECT slug FROM redirect WHERE url = ?", requestUrl).Scan(&slug) - if err == nil { - writeShortenedUrl(w, slug) - return + slug := r.URL.Query().Get("slug") + if slug == "" { + _ = db.QueryRow("SELECT slug FROM redirect WHERE url = ?", requestUrl).Scan(&slug) } - var exists = true - for exists == true { - slug = generateSlug() - err, exists = slugExists(slug) - if err != nil { - http.Error(w, err.Error(), http.StatusInternalServerError) + if slug != "" { + if _, e := slugExists(slug); e { + writeShortenedUrl(w, slug) return } + } else { + var exists = true + for exists == true { + slug = generateSlug() + var err error + err, exists = slugExists(slug) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + } } stmt, err := db.Prepare("INSERT INTO redirect (slug, url, hits) VALUES (?, ?, ?)")