Merge pull request #22 from bat-bs/reporting

chore: fix to many open DB Connections

Author: hofq

.env.example

@@ -6,4 +6,6 @@ API_VERSION=2024-02-15-preview
 RESSOURCE_NAME=gpt-35-turbo
 BASE_URL=openai.azure.com
 
-DEFAULT_BACKEND=azure
+DEFAULT_BACKEND=azure
+
+DATABASE_PATH=postgresql://user:password@server:5432/db?sslmode=disable

api/apiserver.go

@@ -11,13 +11,13 @@ import (
 	"github.com/Masterminds/sprig/v3"
 )
 
-func ApiInit(mux *http.ServeMux, a *auth.Auth) {
+func ApiInit(mux *http.ServeMux, a *auth.Auth, db *db.Database) {
 	timeZone, ok := os.LookupEnv("TIMEZONE")
 	if !ok {
 		timeZone = "Europe/Berlin"
 	}
 
-	api := ApiHandler{db.NewDB(), a, timeZone}
+	api := ApiHandler{db, a, timeZone}
 	mux.HandleFunc("/api2/user/widget", api.GetUserWidget)
 	mux.HandleFunc("/api2/user/logout", api.LogoutUser)
 	mux.HandleFunc("/api2/admin/table/get", api.GetAdminTable)

apiproxy/proxy.go

@@ -9,6 +9,7 @@ import (
 	"net/http"
 	"net/http/httputil"
 	"net/url"
+	db "openai-api-proxy/db"
 	"os"
 	"strings"
 )
@@ -42,7 +43,7 @@ var (
 	backendProxy = make(map[string]*httputil.ReverseProxy)
 )
 
-func Init(mux *http.ServeMux) {
+func Init(mux *http.ServeMux, db *db.Database) {
 	// Setup Azure Vars and Connection String
 	azconf := &AzureConfig{
 		DeploymentName: os.Getenv("DEPLOYMENT_NAME"),
@@ -51,14 +52,21 @@ func Init(mux *http.ServeMux) {
 		BaseUrl:        os.Getenv("BASE_URL"),
 	}
 	defaultBackend = os.Getenv("DEFAULT_BACKEND")
-
-	h := &baseHandle{azconf}
+	rc := &ResponseConf{
+		db: db,
+	}
+	h := &baseHandle{
+		db: db,
+		az: azconf,
+		rc: rc}
 	mux.Handle("/api/", h)
 
 }
 
 type baseHandle struct {
+	db *db.Database
 	az *AzureConfig
+	rc *ResponseConf
 }
 
 func (h *baseHandle) ServeHTTP(w http.ResponseWriter, r *http.Request) {
@@ -122,7 +130,7 @@ func (h *baseHandle) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
 // Since OpenAI and Azure API are not really compatible, we need 2 different handler functions
 func (h *baseHandle) HandleAzure(w http.ResponseWriter, r *http.Request, backend string) {
-	azureToken := ValidateToken(w, r)
+	azureToken := h.ValidateToken(w, r)
 	if azureToken == "" {
 		http.Error(w, "Error Processing Request", http.StatusUnauthorized)
 		return
@@ -150,7 +158,7 @@ func (h *baseHandle) HandleAzure(w http.ResponseWriter, r *http.Request, backend
 	log.Printf("Proxying request to Azure backend: %s", actualURL.String())
 	r.Body.Close()
 
-	proxy.ModifyResponse = NewResponse
+	proxy.ModifyResponse = h.rc.NewResponse
 
 	proxy.ServeHTTP(w, r)
 

apiproxy/response.go

@@ -11,9 +11,13 @@ import (
 	"strings"
 )
 
+type ResponseConf struct {
+	db *db.Database
+}
+
 type Response struct {
 	rs       *http.Response
-	db       *db.Database
+	rc       *ResponseConf
 	apiKeyID string
 	content  Content
 }
@@ -27,18 +31,18 @@ type Content struct {
 	} `json:"usage"`
 }
 
-func NewResponse(in *http.Response) error {
+func (rc *ResponseConf) NewResponse(in *http.Response) error {
 
 	r := Response{
 		rs: in,
-		db: db.NewDB(),
+		rc: rc,
 	}
 
 	err := r.ReadValues()
 	if err != nil {
 		return err
 	}
-	go r.ProcessValues()
+	r.ProcessValues()
 	return nil
 }
 
@@ -47,7 +51,7 @@ func (r *Response) GetApiKeyUUID() string {
 
 	apiKey := strings.TrimPrefix(header, "Bearer ")
 
-	hashes, err := r.db.LookupApiKeys("*")
+	hashes, err := r.rc.db.LookupApiKeys("*")
 	if err != nil || len(hashes) == 0 {
 		log.Println("Error while requesting API Keys from DB", err)
 	}
@@ -69,7 +73,6 @@ func (r *Response) ReadValues() error {
 }
 
 func (r *Response) ProcessValues() {
-	defer r.db.Close()
 	c := r.content
 	if c.Object != "chat.completion" {
 		log.Printf("Untested API Endpoint '%s' is used, check the Request in the DB: %s", c.Object, c.ID)
@@ -86,7 +89,7 @@ func (r *Response) ProcessValues() {
 		TokenCountComplete: c.Usage.CompletionTokens,
 		Model:              c.Model,
 	}
-	if err := r.db.WriteRequest(&rq); err != nil {
+	if err := r.rc.db.WriteRequest(&rq); err != nil {
 		log.Println("error processing Response: Response could not be written to db DB")
 		log.Println(err)
 		return

apiproxy/validateToken.go

@@ -25,7 +25,7 @@ func CompareToken(hashes []db.ApiKey, apiKey string) (string, error) {
 	return "", err
 }
 
-func ValidateToken(w http.ResponseWriter, r *http.Request) string {
+func (h *baseHandle) ValidateToken(w http.ResponseWriter, r *http.Request) string {
 	header := r.Header.Get(authHeader)
 
 	apiKey := strings.TrimPrefix(header, "Bearer ")
@@ -36,9 +36,7 @@ func ValidateToken(w http.ResponseWriter, r *http.Request) string {
 		return ""
 	}
 
-	db := db.NewDB()
-
-	hashes, err := db.LookupApiKeys("*")
+	hashes, err := h.db.LookupApiKeys("*")
 	if err != nil || len(hashes) == 0 {
 		log.Println("Error while requesting API Keys from DB", err)
 		http.Error(w, "401 - Token Invalid", http.StatusUnauthorized)

auth/auth.go

@@ -17,7 +17,7 @@ import (
 	"golang.org/x/oauth2"
 )
 
-func Init(mux *http.ServeMux) (a *Auth) {
+func Init(mux *http.ServeMux, db *db.Database) (a *Auth) {
 	ctx := context.Background()
 	issuer, ok := os.LookupEnv("ISSUER")
 	if !ok {
@@ -63,8 +63,10 @@ func Init(mux *http.ServeMux) (a *Auth) {
 	}
 
 	var verifier = provider.Verifier(&oidc.Config{ClientID: clientId})
+
 	a = &Auth{
 		oauth2Config: oauth2Config,
+		db:           db,
 		ctx:          context.Background(),
 		verifier:     verifier,
 		provider:     provider,
@@ -89,6 +91,7 @@ type Claims struct {
 
 type Auth struct {
 	oauth2Config *oauth2.Config
+	db           *db.Database
 	ctx          context.Context
 	verifier     *oidc.IDTokenVerifier
 	provider     *oidc.Provider
@@ -205,9 +208,8 @@ func (a *Auth) ValidateAdminSession(w http.ResponseWriter, r *http.Request) (boo
 		log.Println("Claims not found")
 		return false, err
 	}
-	db := db.NewDB()
-	defer db.Close()
-	user, err := db.GetUser(claims.Sub)
+
+	user, err := a.db.GetUser(claims.Sub)
 	if err != nil {
 		return false, err
 	}

cmd/main.go

@@ -8,6 +8,9 @@ import (
 	auth "openai-api-proxy/auth"
 	db "openai-api-proxy/db"
 	web "openai-api-proxy/webui"
+	"os"
+	"os/signal"
+	"syscall"
 
 	"github.com/joho/godotenv"
 )
@@ -18,15 +21,35 @@ func main() {
 	if err != nil {
 		log.Println("Warning: not able to loading Env File", err)
 	}
-	db.DatabaseInit()
+	db := db.DatabaseInit()
+	defer db.Close()
+	osExit(db)
+	defer log.Println("Closing DB Clients :)")
+
 	mux := http.NewServeMux()
-	a := auth.Init(mux)
+	a := auth.Init(mux, db)
 	//
-	proxy.Init(mux)        // Start AI Proxy
-	go web.Init(mux, a)    // Start Web UI
-	go api.ApiInit(mux, a) // Start Backend API
+	proxy.Init(mux, db)        // Start AI Proxy
+	go web.Init(mux, a)        // Start Web UI
+	go api.ApiInit(mux, a, db) // Start Backend API
 
 	log.Printf("Serving on http://localhost:%d", 8082)
 	log.Fatal(http.ListenAndServe(":8082", mux))
 
 }
+
+// Close DB on Program Exit
+func osExit(db *db.Database) {
+	sigc := make(chan os.Signal, 1)
+	signal.Notify(sigc,
+		syscall.SIGHUP,
+		syscall.SIGINT,
+		syscall.SIGTERM,
+		syscall.SIGQUIT)
+	go func() {
+		s := <-sigc
+		log.Printf("Exit: %s Closing DB Clients :)", s)
+		db.Close()
+		os.Exit(1)
+	}()
+}

db/database.go

@@ -24,18 +24,18 @@ type ApiKey struct {
 	TokenCountComplete *int
 }
 
-func DatabaseInit() {
+func DatabaseInit() *Database {
 	createTable, err := os.ReadFile("db/schema.sql")
 	if err != nil {
 		log.Fatal("cannot load schema file: ", err)
 	}
 
 	d := NewDB()
-	defer d.Close()
 	d.Migrate()
 	if _, err := d.db.Exec(string(createTable)); err != nil {
 		log.Fatal(err)
 	}
+	return d
 
 }