storage.go

// Copyright © Paul Tötterman <paul.totterman@gmail.com>. All rights reserved.

package main

import (
	"context"
	"database/sql"
	"fmt"
	"net"
	"strconv"
)

type execer interface {
	ExecContext(ctx context.Context, query string, args ...any) (
		sql.Result, error)
}

func ensureSchema(e execer) error {
	const q = `
CREATE TABLE IF NOT EXISTS urls (
    created timestamp with time zone NOT NULL DEFAULT now(),
    id bigint GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,
    hits bigint NOT NULL DEFAULT 0,
    name text NOT NULL UNIQUE,
    url text NOT NULL,
    "user" text NOT NULL
);

CREATE TABLE IF NOT EXISTS hits (
    created timestamp with time zone NOT NULL DEFAULT now(),
    url_id bigint NOT NULL REFERENCES urls (id) ON DELETE CASCADE,
    remotehost inet,
    referrer text,
    agent text
);
`

	ctx := context.Background()

	if _, err := e.ExecContext(ctx, q); err != nil {
		return fmt.Errorf("failed querying DB: %w", err)
	}

	return nil
}

// newPostgresDB returns an initialized postgresDB.
func newPostgresDB() (*sql.DB, error) {
	db, err := sql.Open("postgres", conf.DB)
	if err != nil {
		return nil, fmt.Errorf("failed opening DB: %w", err)
	}

	if err := ensureSchema(db); err != nil {
		return nil, fmt.Errorf("failed ensuring schema: %w", err)
	}

	return db, nil
}

// getURLnID returns URL and its ID.
func getURLnID(ctx context.Context, tx *sql.Tx, name string) (string, int64,
	error,
) {
	const q = `
UPDATE
    urls
SET
    hits = hits + 1
WHERE
    name = $1
RETURNING
    id,
    url;
`

	var (
		id  int64
		url string
	)

	//nolint:execinquery
	if err := tx.QueryRowContext(ctx, q, name).Scan(&id, &url); err != nil {
		return "", 0, fmt.Errorf("failed querying DB: %w", err)
	}

	return url, id, nil
}

// getIDnUser returns the URL's ID and user.
//
//nolint:unparam
func getIDnUser(ctx context.Context, tx *sql.Tx, name string) (int64, string,
	error,
) {
	const q = `
SELECT
    id,
    "user"
FROM
    urls
WHERE
    name = $1;
`

	var (
		id   int64
		user string
	)

	if err := tx.QueryRowContext(ctx, q, name).Scan(&id,
		&user); err != nil {
		return 0, "", fmt.Errorf("failed querying DB: %w", err)
	}

	return id, user, nil
}

// removeURL removes the URL speficied.
func removeURL(ctx context.Context, tx *sql.Tx, name string) error {
	const q = `
DELETE FROM urls
WHERE name = $1;
`

	if _, err := tx.ExecContext(ctx, q, name); err != nil {
		return fmt.Errorf("failed querying DB: %w", err)
	}

	return nil
}

// addHit adds a hit to the specific URL.
func addHit(ctx context.Context, tx *sql.Tx, urlID int64, ip net.IP,
	agent string, referrer *string,
) error {
	const q = `
INSERT INTO hits (
    url_id,
    remotehost,
    agent,
    referrer)
VALUES (
    $1,
    $2,
    $3,
    $4);
`

	if _, err := tx.ExecContext(ctx, q, urlID, ip.String(), agent,
		referrer); err != nil {
		return fmt.Errorf("failed querying DB: %w", err)
	}

	return nil
}

// addURL adds a new URL to the database.
func addURL(ctx context.Context, tx *sql.Tx, name, url, user string) error {
	const q = `
INSERT INTO urls (
    name,
    url,
    "user")
VALUES (
    $1,
    $2,
    $3);
`

	if _, err := tx.ExecContext(ctx, q, name, url, user); err != nil {
		return fmt.Errorf("failed querying DB: %w", err)
	}

	return nil
}

// urlsForUser returns all URLs for the given user.
func urlsForUser(ctx context.Context, tx *sql.Tx, user string) (
	[]map[string]string, error,
) {
	const q = `
SELECT
    name,
    url,
    hits
FROM
    urls
WHERE
    "user" = $1;
`

	//nolint:sqlclosecheck
	rows, err := tx.QueryContext(ctx, q, user)
	if err != nil {
		return nil, fmt.Errorf("failed querying DB: %w", err)
	}

	defer func(rows *sql.Rows) {
		if err = rows.Close(); err != nil {
			panic(err)
		}
	}(rows)

	urls := []map[string]string{}

	for rows.Next() {
		var (
			name, url string
			hits      int
		)

		if err = rows.Scan(&name, &url, &hits); err != nil {
			return nil, fmt.Errorf("failed querying DB: %w", err)
		}

		urls = append(urls, map[string]string{
			"name": name,
			"url":  url,
			"hits": strconv.Itoa(hits),
		})
	}

	err = rows.Err()
	if err != nil {
		return nil, fmt.Errorf("failed querying DB: %w", err)
	}

	return urls, nil
}