Build an OTP-Based Authentication Server with Go: Part File Uploads and Graceful Shutdown

This installment refines our Go authentication server by adding file upload capabilities, streamlining development with a Makefile, and implementing graceful server shutdown. This prevents abrupt termination and ensures all in-progress tasks complete before the server stops.

Graceful Shutdown Implementation

A new cmd/api/server.go file centralizes server management. A goroutine monitors termination signals (SIGINT, SIGTERM). Upon receiving a signal, it gracefully shuts down the server. Comments clarify each step.

package main

import (
    "context"
    "errors"
    "fmt"
    "net/http"
    "os"
    "os/signal"
    "syscall"
    "time"
)

// ... (Existing application code) ...

func (app *application) serve(router http.Handler) error {
    // Server initialization with timeouts for idle, read, and write operations.
    srv := &http.Server{
        Addr:         fmt.Sprintf(":%d", app.config.port),
        Handler:      app.recoverPanic(app.authenticate(router)),
        IdleTimeout:  time.Minute,
        ReadTimeout:  10 * time.Second,
        WriteTimeout: 30 * time.Second,
    }

    // Channel to manage errors during shutdown.
    shutdownError := make(chan error)

    // Goroutine to listen for termination signals and initiate graceful shutdown.
    go func() {
        quit := make(chan os.Signal, 1)
        signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
        s := <-quit // Wait for signal
        fmt.Println("Shutting down server...")

        // Context with timeout for graceful shutdown.
        ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
        defer cancel()

        // Attempt graceful shutdown.
        shutdownError <- srv.Shutdown(ctx)
    }()

    // Start the server and report any errors.
    fmt.Printf("Starting server on port %d\n", app.config.port)
    err := srv.ListenAndServe()
    if !errors.Is(err, http.ErrServerClosed) {
        return fmt.Errorf("server error: %w", err)
    }
    return <-shutdownError
}

// ... (Rest of application code) ...

In main.go, replace the srv.ListenAndServe() block with err = app.serve(router); if err != nil { logger.Fatal(err, nil) }. Testing this involves using the retry mechanism (remove credentials for continuous retries) and interrupting the server with Ctrl C; it should wait for retries to finish before shutting down.

Automating Tasks with Makefile

A Makefile automates repetitive commands. Add the following to your project's Makefile:

# Include variables from the .envrc file
include .envrc

## help: print this help message
.PHONY: help
help:
    @echo 'Usage:'
    @sed -n 's/^##//p' ${MAKEFILE_LIST} | column -t -s ':' | sed -e 's/^/ /'

.PHONY: confirm
confirm:
    @echo -n 'Are you sure? [y/N] ' && read ans && [ $${ans:-N} = y ]

## run/api: run the cmd/api application
.PHONY: run/api
run/api:
    go run ./cmd/api

## db/psql: connect to the database using psql
.PHONY: db/psql
db/psql:
    psql ${GREENLIGHT_DB_DSN}

## db/migrations/new name=: create a new database migration
.PHONY: db/migrations/new
db/migrations/new:
    @echo "Migrating ${name}"
    migrate create -seq -ext=.sql -dir=./migrations ${name}

## db/migrations/up: apply all up database migrations
.PHONY: db/migrations/up
db/migrations/up: confirm
    @echo "Running migrations"
    migrate -path=./migrations -database=${GREENLIGHT_DB_DSN} up

## audit: tidy and vendor dependencies and format, vet and test all code
.PHONY: audit
audit: vendor
    @echo 'Formatting code...'
    go fmt ./...
    @echo 'Vetting code...'
    go vet ./...
    staticcheck ./...
    @echo 'Running tests...'
    go test -race -vet=off ./...

## vendor: tidy and vendor dependencies
.PHONY: vendor
vendor:
    @echo 'Tidying and verifying module dependencies...'
    go mod tidy
    go mod verify
    @echo 'Vendoring dependencies...'
    go mod vendor

## build/api: build the cmd/api application
.PHONY: build/api
build/api:
    @echo 'Building cmd/api...'
    go build -o=./bin/api ./cmd/api
    GOOS=linux GOARCH=amd64 go build -o=./bin/linux_amd64/api ./cmd/api

Run commands like make help, make run/api, make db/migrations/new name=my_migration, make db/migrations/up, make audit, or make build/api. Additional commands can be added following this structure.

File Upload and Tracking

A new database table tracks uploaded files:

Create migration make db/migrations/new name=create-creatives.

000003_create-creatives.up.sql:

CREATE TABLE IF NOT EXISTS creatives (
    id bigserial PRIMARY KEY,
    user_id bigint NOT NULL REFERENCES users ON DELETE CASCADE,
    creative_url text NOT NULL,
    scheduled_at DATE NOT NULL,
    created_at timestamp(0) with time zone NOT NULL DEFAULT NOW()
);

000003_create-creatives.down.sql:

DROP TABLE IF EXISTS creatives;

internal/data/creatives.go:

package data

import (
    "context"
    "database/sql"
    "time"

    "github.com/lib/pq"
)

// ... (Creative struct and CreativeModel struct) ...

func (c *CreativeModel) Insert(creative *Creative) error {
    query := `INSERT INTO creatives (user_id, creative_url, scheduled_at)
            VALUES (, , )
            RETURNING id, created_at`

    ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
    defer cancel()

    args := []interface{}{creative.UserID, creative.CreativeURL, creative.ScheduledAt}
    return c.DB.QueryRowContext(ctx, query, args...).Scan(&creative.ID, &creative.CreatedAt)
}

func (c *CreativeModel) GetScheduledCreatives() (map[string][]Creative, error) {
    query := `
        SELECT id, user_id, creative_url, scheduled_at, created_at 
        FROM creatives 
        WHERE scheduled_at = ANY()
    `

    dates := []time.Time{
        time.Now().Truncate(24 * time.Hour),
        time.Now().AddDate(0, 0, 1).Truncate(24 * time.Hour),
    }

    ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
    defer cancel()

    rows, err := c.DB.QueryContext(ctx, query, pq.Array(dates))
    if err != nil {
        return nil, err
    }
    defer rows.Close()

    creatives := map[string][]Creative{
        "today":    {},
        "tomorrow": {},
    }

    for rows.Next() {
        var creative Creative
        err := rows.Scan(&creative.ID, &creative.UserID, &creative.CreativeURL, &creative.ScheduledAt, &creative.CreatedAt)
        if err != nil {
            return nil, err
        }
        if creative.ScheduledAt.Equal(dates[0]) {
            creatives["today"] = append(creatives["today"], creative)
        } else if creative.ScheduledAt.Equal(dates[1]) {
            creatives["tomorrow"] = append(creatives["tomorrow"], creative)
        }
    }

    return creatives, rows.Err()
}

cmd/api/creatives.go: (Handles file uploads and retrieval of scheduled creatives)

package main

import (
    "fmt"
    "io"
    "net/http"
    "os"
    "path"
    "strings"
    "time"

    "github.com/google/uuid"
    "github.com/vishaaxl/cheershare/internal/data"
)

const MaxFileSize = 10 << 20 // 10MB

// ... (Existing code) ...

func (app *application) uploadCreativeHandler(w http.ResponseWriter, r *http.Request) {
    // ... (File upload handling logic) ...
}

func (app *application) getScheduledCreativesHandler(w http.ResponseWriter, r *http.Request) {
    // ... (Retrieve scheduled creatives logic) ...
}

Add the following to your main.go or server initialization to create the uploads directory if it doesn't exist:

uploadDir := "./uploads"
if _, err := os.Stat(uploadDir); os.IsNotExist(err) {
    err := os.MkdirAll(uploadDir, os.ModePerm)
    if err != nil {
        fmt.Println("Unable to create uploads directory:", err)
    }
}

router.HandlerFunc(http.MethodPost, "/upload-creative", app.requireAuthenticatedUser(app.uploadCreativeHandler))
router.HandlerFunc(http.MethodGet, "/scheduled", app.requireAuthenticatedUser(app.getScheduledCreativesHandler))

Remember to replace placeholder comments with actual implementation details for file handling and error management within the handlers. This completes the core functionality, excluding payment processing (to be covered in future updates). The provided link to Part 3 is maintained.

The above is the detailed content of Build an OTP-Based Authentication Server with Go: Part File Uploads and Graceful Shutdown. For more information, please follow other related articles on the PHP Chinese website!