| 1 |
package main |
| 2 |
|
| 3 |
import ( |
| 4 |
"github.com/jackc/pgx" |
| 5 |
log "gopkg.in/inconshreveable/log15.v2" |
| 6 |
"io/ioutil" |
| 7 |
"net/http" |
| 8 |
"os" |
| 9 |
) |
| 10 |
|
| 11 |
var pool *pgx.ConnPool |
| 12 |
|
| 13 |
// afterConnect creates the prepared statements that this application uses |
| 14 |
func afterConnect(conn *pgx.Conn) (err error) { |
| 15 |
_, err = conn.Prepare("getUrl", ` |
| 16 |
select url from shortened_urls where id=$1 |
| 17 |
`) |
| 18 |
if err != nil { |
| 19 |
return |
| 20 |
} |
| 21 |
|
| 22 |
_, err = conn.Prepare("deleteUrl", ` |
| 23 |
delete from shortened_urls where id=$1 |
| 24 |
`) |
| 25 |
if err != nil { |
| 26 |
return |
| 27 |
} |
| 28 |
|
| 29 |
// There technically is a small race condition in doing an upsert with a CTE |
| 30 |
// where one of two simultaneous requests to the shortened URL would fail |
| 31 |
// with a unique index violation. As the point of this demo is pgx usage and |
| 32 |
// not how to perfectly upsert in PostgreSQL it is deemed acceptable. |
| 33 |
_, err = conn.Prepare("putUrl", ` |
| 34 |
with upsert as ( |
| 35 |
update shortened_urls |
| 36 |
set url=$2 |
| 37 |
where id=$1 |
| 38 |
returning * |
| 39 |
) |
| 40 |
insert into shortened_urls(id, url) |
| 41 |
select $1, $2 where not exists(select 1 from upsert) |
| 42 |
`) |
| 43 |
return |
| 44 |
} |
| 45 |
|
| 46 |
func getUrlHandler(w http.ResponseWriter, req *http.Request) { |
| 47 |
var url string |
| 48 |
err := pool.QueryRow("getUrl", req.URL.Path).Scan(&url) |
| 49 |
switch err { |
| 50 |
case nil: |
| 51 |
http.Redirect(w, req, url, http.StatusSeeOther) |
| 52 |
case pgx.ErrNoRows: |
| 53 |
http.NotFound(w, req) |
| 54 |
default: |
| 55 |
http.Error(w, "Internal server error", http.StatusInternalServerError) |
| 56 |
} |
| 57 |
} |
| 58 |
|
| 59 |
func putUrlHandler(w http.ResponseWriter, req *http.Request) { |
| 60 |
id := req.URL.Path |
| 61 |
var url string |
| 62 |
if body, err := ioutil.ReadAll(req.Body); err == nil { |
| 63 |
url = string(body) |
| 64 |
} else { |
| 65 |
http.Error(w, "Internal server error", http.StatusInternalServerError) |
| 66 |
return |
| 67 |
} |
| 68 |
|
| 69 |
if _, err := pool.Exec("putUrl", id, url); err == nil { |
| 70 |
w.WriteHeader(http.StatusOK) |
| 71 |
} else { |
| 72 |
http.Error(w, "Internal server error", http.StatusInternalServerError) |
| 73 |
} |
| 74 |
} |
| 75 |
|
| 76 |
func deleteUrlHandler(w http.ResponseWriter, req *http.Request) { |
| 77 |
if _, err := pool.Exec("deleteUrl", req.URL.Path); err == nil { |
| 78 |
w.WriteHeader(http.StatusOK) |
| 79 |
} else { |
| 80 |
http.Error(w, "Internal server error", http.StatusInternalServerError) |
| 81 |
} |
| 82 |
} |
| 83 |
|
| 84 |
func urlHandler(w http.ResponseWriter, req *http.Request) { |
| 85 |
switch req.Method { |
| 86 |
case "GET": |
| 87 |
getUrlHandler(w, req) |
| 88 |
|
| 89 |
case "PUT": |
| 90 |
putUrlHandler(w, req) |
| 91 |
|
| 92 |
case "DELETE": |
| 93 |
deleteUrlHandler(w, req) |
| 94 |
|
| 95 |
default: |
| 96 |
w.Header().Add("Allow", "GET, PUT, DELETE") |
| 97 |
w.WriteHeader(http.StatusMethodNotAllowed) |
| 98 |
} |
| 99 |
} |
| 100 |
|
| 101 |
func main() { |
| 102 |
var err error |
| 103 |
connPoolConfig := pgx.ConnPoolConfig{ |
| 104 |
ConnConfig: pgx.ConnConfig{ |
| 105 |
Host: "127.0.0.1", |
| 106 |
User: "jack", |
| 107 |
Password: "jack", |
| 108 |
Database: "url_shortener", |
| 109 |
Logger: log.New("module", "pgx"), |
| 110 |
}, |
| 111 |
MaxConnections: 5, |
| 112 |
AfterConnect: afterConnect, |
| 113 |
} |
| 114 |
pool, err = pgx.NewConnPool(connPoolConfig) |
| 115 |
if err != nil { |
| 116 |
log.Crit("Unable to create connection pool", "error", err) |
| 117 |
os.Exit(1) |
| 118 |
} |
| 119 |
|
| 120 |
http.HandleFunc("/", urlHandler) |
| 121 |
|
| 122 |
log.Info("Starting URL shortener on localhost:8080") |
| 123 |
err = http.ListenAndServe("localhost:8080", nil) |
| 124 |
if err != nil { |
| 125 |
log.Crit("Unable to start web server", "error", err) |
| 126 |
os.Exit(1) |
| 127 |
} |
| 128 |
} |
