StateDB: The first few days, Part D - Reorganizing the code

In Part C, I was trying inserts, deletes, and updates. The next step was reorganizing the code and moving things into functions.

I started separating the database connection from the code that works with test items. main would read the configuration, call those functions, and print the results.

Moving the connection into a package

The connection code became part of a database package:

package database

import (
	"context"
	"fmt"

	"github.com/jackc/pgx/v5/pgxpool"
)

type Config struct {
	User     string
	Password string
	Host     string
	Port     string
	Database string
}

func Connect(ctx context.Context, config Config) (*pgxpool.Pool, error) {
	connectionString := fmt.Sprintf(
		"postgres://%s:%s@%s:%s/%s",
		config.User,
		config.Password,
		config.Host,
		config.Port,
		config.Database,
	)
	return pgxpool.New(ctx, connectionString)
}

The configuration type is now database.Config, and the function is database.Connect. The capital letters make those names exported, so another package can use them.

This function receives the settings instead of reading environment variables itself. That part stays in main.

Something I learned while writing this post: building the connection URL with fmt.Sprintf doesn’t escape special characters in credentials. Using net/url with url.UserPassword handles that escaping. I didn’t know that when I wrote this function (and still haven’t actually tested), so the code above is still my initial version.

Creating the pool doesn’t confirm a connection yet, so I still call Ping in main.

Giving a test item a type

The query code moved into a testitem package, with a struct for a row:

package testitem

import (
	"context"
	"log"
	"time"

	"github.com/jackc/pgx/v5/pgxpool"
)

type TestItem struct {
	ID        int
	Name      string
	Status    string
	CreatedAt time.Time
}

Before this, I was scanning columns into separate variables. Now I could return one TestItem containing the values that belong together.

Getting one item

The first query function takes an ID and returns a TestItem and an error:

func GetByID(
	ctx context.Context,
	db *pgxpool.Pool,
	id int,
) (TestItem, error) {
	var item TestItem

	err := db.QueryRow(
		ctx,
		`
		select id, name, status, created_at
		from test_items
		where id = $1
		`,
		id,
	).Scan(
		&item.ID,
		&item.Name,
		&item.Status,
		&item.CreatedAt,
	)

	return item, err
}

QueryRow and Scan are chained together here. The error comes from Scan, and the column values go into the fields of item through pointers.

The function returns the item and error to its caller instead of printing the row. If no row matches, Scan returns pgx.ErrNoRows, and the caller needs to handle that before using the item.

Getting all items

The next function returns []TestItem, a slice of items:

func GetAll(
	ctx context.Context,
	db *pgxpool.Pool,
) ([]TestItem, error) {
	rows, err := db.Query(
		ctx,
		`
		select id, name, status, created_at
		from test_items
		`,
	)
	if err != nil {
		log.Fatal(err)
	}
	defer rows.Close()

	var items []TestItem

	for rows.Next() {
		var item TestItem

		if err := rows.Scan(
			&item.ID,
			&item.Name,
			&item.Status,
			&item.CreatedAt,
		); err != nil {
			log.Fatal(err)
		}

		items = append(items, item)
	}

	if err := rows.Err(); err != nil {
		return nil, err
	}

	return items, nil
}

The loop still reads one row at a time, but now append adds each item to a slice. The caller gets the collection back and decides what to do with it. If there are no rows and no error, this returns a nil slice and a nil error.

Calling the functions from main

In main, I import the two packages:

import (
	"statedb/internal/database"
	"statedb/internal/testitem"
)

These are the project imports alongside context, fmt, log, and os. The statedb prefix needs to match the module path in go.mod.

The configuration and connection setup now use the database package:

ctx := context.Background()

config := database.Config{
	User:     os.Getenv("DB_USER_NAME"),
	Password: os.Getenv("DB_PASSWORD"),
	Host:     os.Getenv("DB_HOST"),
	Port:     os.Getenv("DB_PORT"),
	Database: os.Getenv("DB_NAME"),
}

if config.User == "" ||
	config.Password == "" ||
	config.Host == "" ||
	config.Port == "" ||
	config.Database == "" {
	log.Fatal("missing required database configuration")
}

fmt.Println("Connecting to StateDB")

db, err := database.Connect(ctx, config)
if err != nil {
	log.Fatal(err)
}
defer db.Close()

if err := db.Ping(ctx); err != nil {
	log.Fatal(err)
}

fmt.Println("Connected to StateDB")

In this version, the username comes from DB_USER_NAME, rather than the DB_USER used in the earlier example.

Then I can ask for an item without putting its SQL in main:

itemFound, err := testitem.GetByID(ctx, db, 4)
if err != nil {
	log.Fatal(err)
} else {
	fmt.Printf(
		"ID: %d, Name: %s, Status: %s, Created At: %v\n",
		itemFound.ID,
		itemFound.Name,
		itemFound.Status,
		itemFound.CreatedAt,
	)
}

I also tried ID 20 with log.Println(err) to check what happens when a value i pass doesn’t exist.

Reading everything uses the other function:

items, err := testitem.GetAll(ctx, db)
if err != nil {
	log.Fatal(err)
}
for _, item := range items {
	fmt.Printf(
		"ID: %d, Name: %s, Status: %s, Created At: %v\n",
		item.ID,
		item.Name,
		item.Status,
		item.CreatedAt,
	)
}

range gives me an index and an item from the slice. I only need the item here, so _ discards the index.

The log.Fatal calls in main still exit immediately, including skipping defer db.Close() on those error paths. Moving the work into a function that returns an error would allow its deferred cleanup to run before main decides to exit.

Still a work in progress

The insert, delete, and update calls were still in main at this point. I was inserting item-8, deleting by that name, and updating item-seven to item-7. I hadn’t moved those operations into the package yet.

But the connection setup and read queries now had functions I could call. I could look at main and see the steps of the program, then go into the other packages when I needed to work on the details.

That’s where I got to with this reorganization.