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