Browse Source

feat: mysql support

mysql
evanchen333 6 months ago
parent
commit
02bf0d68d5
  1. 10
      Makefile
  2. 28
      example/main.go
  3. 2
      go.mod
  4. 11
      go.sum
  5. 15
      limiter.go

10
Makefile

@ -1,2 +1,10 @@
run:
go run example/main.go
go run example/main.go
start-mysql:
docker run --name limiter_db --rm -d \
-e MYSQL_ROOT_PASSWORD=admin \
-e MYSQL_DATABASE=limiter \
-e MYSQL_ROOT_HOST=% \
-p 3306:3306 \
mysql --default-authentication-plugin=mysql_native_password

28
example/main.go

@ -1,9 +1,13 @@
package main
import (
"database/sql"
"fmt"
"net/http"
"os"
"time"
_ "github.com/go-sql-driver/mysql"
limiter "github.com/mutsuki333/rate-limiter"
)
@ -19,9 +23,27 @@ func handler(w http.ResponseWriter, r *http.Request) {
}
func main() {
l = limiter.Default()
// l.Interval = time.Minute * 10
l.Limit = 5
if len(os.Args) > 1 {
var db *sql.DB
var err error
if os.Args[1] == "mysql" {
db, err = sql.Open("mysql", "root:admin@/limiter?charset=utf8")
} else {
db, err = sql.Open("sqlite3", os.Args[1]+".db")
}
if err != nil {
panic(err)
}
l = &limiter.Limiter{
Interval: time.Minute,
Limit: 5,
Store: db,
}
l.Init()
} else {
l = limiter.Default()
}
http.HandleFunc("/hit", handler)
http.ListenAndServe(":8080", nil)

2
go.mod

@ -3,6 +3,6 @@ module github.com/mutsuki333/rate-limiter
go 1.15
require (
github.com/go-sql-driver/mysql v1.5.0
github.com/mattn/go-sqlite3 v1.14.6
gorm.io/driver/sqlite v1.1.4
)

11
go.sum

@ -1,11 +1,4 @@
github.com/jinzhu/inflection v1.0.0 h1:K317FqzuhWc8YvSVlFMCCUb36O/S9MCKRDI7QkRKD/E=
github.com/jinzhu/inflection v1.0.0/go.mod h1:h+uFLlag+Qp1Va5pdKtLDYj+kHp5pxUVkryuEj+Srlc=
github.com/jinzhu/now v1.1.1 h1:g39TucaRWyV3dwDO++eEc6qf8TVIQ/Da48WmqjZ3i7E=
github.com/jinzhu/now v1.1.1/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/z8=
github.com/mattn/go-sqlite3 v1.14.5/go.mod h1:WVKg1VTActs4Qso6iwGbiFih2UIHo0ENGwNd0Lj+XmI=
github.com/go-sql-driver/mysql v1.5.0 h1:ozyZYNQW3x3HtqT1jira07DN2PArx2v7/mN66gGcHOs=
github.com/go-sql-driver/mysql v1.5.0/go.mod h1:DCzpHaOWr8IXmIStZouvnhqoel9Qv2LBy8hT2VhHyBg=
github.com/mattn/go-sqlite3 v1.14.6 h1:dNPt6NO46WmLVt2DLNpwczCmdV5boIZ6g/tlDrlRUbg=
github.com/mattn/go-sqlite3 v1.14.6/go.mod h1:NyWgC/yNuGj7Q9rpYnZvas74GogHl5/Z4A/KQRfk6bU=
gorm.io/driver/sqlite v1.1.4 h1:PDzwYE+sI6De2+mxAneV9Xs11+ZyKV6oxD3wDGkaNvM=
gorm.io/driver/sqlite v1.1.4/go.mod h1:mJCeTFr7+crvS+TRnWc5Z3UvwxUN1BGBLMrf5LA9DYw=
gorm.io/gorm v1.20.7 h1:rMS4CL3pNmYq1V5/X+nHHjh1Dx6dnf27+Cai5zabo+M=
gorm.io/gorm v1.20.7/go.mod h1:0HFTzE/SqkGTzK6TlDPPQbAYCluiVvhzoA1+aVyzenw=

15
limiter.go

@ -14,8 +14,7 @@ import (
//Default setting of the limiter, which limits 60hits/minute and uses in-memory sqlite db
func Default() *Limiter {
var err error
// db, err := sql.Open("sqlite3", "file::memory:?cache=shared")
db, err := sql.Open("sqlite3", "test.db")
db, err := sql.Open("sqlite3", "file::memory:?cache=shared")
if err != nil {
panic(err)
}
@ -44,7 +43,15 @@ type Limiter struct {
//Init the db table, and start cleanup goroutine
func (l *Limiter) Init() {
l.Store.Exec(`create table hit (ip string, hit_time datetime);`)
_, err := l.Store.Exec(`drop table if exists hit;`)
if err != nil {
panic(err)
}
_, err = l.Store.Exec(`create table hit (ip varchar(20), hit_time datetime);`)
if err != nil {
panic(err)
}
l.Mux = &sync.Mutex{}
go func() {
for {
time.Sleep(l.Interval)
@ -56,7 +63,7 @@ func (l *Limiter) Init() {
//Rate within interval
func (l *Limiter) Rate(ip string) (rate int, err error) {
row := l.Store.QueryRow(
"select COUNT(*) from hit where ip = ? and hit_time > ?",
"select count(*) from hit where ip = ? and hit_time > ?",
ip,
time.Now().Add(-l.Interval),
)

Loading…
Cancel
Save