package main
import (
"database/sql"
"fmt"
"log"
"net/http"
"os"
"strconv"
_ "github.com/lib/pq" // PostgreSQL driver
)
var db *sql.DB // Global database connection
// dbConnection initializes the database connection.
func dbConnection() (*sql.DB, error) {
connStr := os.Getenv("DATABASE_URL") // You should set this env var, e.g., postgres://username:password@localhost/dbname?sslmode=disable
return sql.Open("postgres", connStr)
}
// livenessHandler is used by Kubernetes to check if the app is alive.
func livenessHandler(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusOK)
fmt.Fprintf(w, "OK")
}
// readinessHandler is used by Kubernetes to check if the app is ready to serve traffic.
func readinessHandler(w http.ResponseWriter, r *http.Request) {
// Check if the database connection is alive
err := db.Ping()
if err != nil {
log.Printf("Readiness probe failed: Database connection error: %v", err)
w.WriteHeader(http.StatusServiceUnavailable)
fmt.Fprintf(w, "NOT READY")
return
}
// If the database is reachable, return 200 OK
w.WriteHeader(http.StatusOK)
fmt.Fprintf(w, "OK")
}
// indexHandler serves the HTML form for the query.
func indexHandler(w http.ResponseWriter, r *http.Request) {
tmpl := `
Query Form
Query Form
`
fmt.Fprintf(w, tmpl)
}
// selectHandler handles HTTP requests and executes a SQL query.
func selectHandler(w http.ResponseWriter, r *http.Request) {
// Get the 'param' query parameter
paramStr := r.URL.Query().Get("param")
if paramStr == "" {
http.Error(w, "Missing 'param' query parameter", http.StatusBadRequest)
return
}
// Convert the param to an integer
param, err := strconv.Atoi(paramStr)
if err != nil {
http.Error(w, "Invalid 'param' query parameter. Must be an integer.", http.StatusBadRequest)
return
}
// Prepare the SQL query to prevent SQL injection
query := "SELECT 42 + $1"
// Execute the query with the provided parameter
var result int
err = db.QueryRow(query, param).Scan(&result)
if err != nil {
log.Printf("Failed to execute query: %v", err)
http.Error(w, "Internal Server Error", http.StatusInternalServerError)
return
}
// Return the result in a simple HTML response
fmt.Fprintf(w, "Result: %d
", result)
}
func main() {
var err error
// Initialize the database connection once at startup
db, err = dbConnection()
if err != nil {
log.Fatalf("Failed to connect to the database: %v", err)
}
defer db.Close()
// Define the handler for the `/` route (serves HTML form)
http.HandleFunc("/", indexHandler)
// Define the handler for the `/query` route (executes SQL query)
http.HandleFunc("/query", selectHandler)
// Define the handler for the `/liveness` probe
http.HandleFunc("/liveness", livenessHandler)
// Define the handler for the `/readiness` probe
http.HandleFunc("/readiness", readinessHandler)
// Start the HTTP server
port := ":8080"
log.Printf("Server starting on port %s\n", port)
err = http.ListenAndServe(port, nil)
if err != nil {
log.Fatalf("Server failed to start: %v", err)
}
}