Version:

Schema Management

This document describes the database schema management system for the pgEdge AI Workbench Collector.

Overview

The Collector uses a migration-based schema management system that provides the following capabilities:

The system automatically creates and updates database schemas at startup.
The system tracks which migrations have been applied.
The system ensures migrations are applied in the correct order.
The system supports idempotent migrations that can run multiple times safely.
The system creates tables, indexes, constraints, and foreign keys.

Architecture

The schema management system consists of several components that work together to maintain the database schema.

SchemaManager

The SchemaManager struct manages all database migrations. The manager maintains a registry of all available migrations, determines which migrations need to be applied, applies pending migrations in order, and tracks migration status in the database.

Migration

Each Migration struct represents a single schema change and contains the following fields:

Version is a unique integer identifying the migration in sequential order.
Description is a human-readable description of the migration.
Up is a function of the form func(pgx.Tx) error that applies the migration on the transaction opened for it.

schema_version Table

The schema_version table tracks which migrations have been applied.

CREATE TABLE schema_version (
    version INTEGER PRIMARY KEY,
    description TEXT NOT NULL,
    applied_at TIMESTAMPTZ NOT NULL
        DEFAULT CURRENT_TIMESTAMP
)

Migration Process

When the Collector starts up, the system executes the following steps:

The Datastore.initializeSchema() method is called.
A new SchemaManager is created with all registered migrations.
The SchemaManager.Migrate() method sorts the migrations by version number, queries the current schema version, applies each pending migration in a transaction, records successful migrations in schema_version, and rolls back on errors.

Current Migrations

The migrations are defined in Go rather than in SQL files. The registerMigrations() method in collector/src/database/schema.go is the authoritative list, and this document does not repeat it, because a second copy drifts as soon as someone adds a migration.

Each entry is a Migration value carrying a Version, a Description and an Up function of the form func(pgx.Tx) error that runs the migration's statements on the transaction the SchemaManager opens for it.

Migration 1 is the consolidated baseline; the Up function creates the complete schema, including every table, index and constraint, along with the seed data for probe configurations and alert rules. Each later migration is an incremental change, such as a new column, a new index or a corrected alert rule, for an installation that already holds the earlier schema. A fresh installation still applies every migration in version order, starting from Migration 1; the later migrations find their changes already present in the baseline, which is why each one must be idempotent.

Versions are allocated sequentially, so the highest Version registered in schema.go is the schema version a current collector converges to. The SchemaManager.LatestVersion() method returns that value at run time, without needing a database connection.

The schema_version table records what has been applied. The Migrate() method inserts a row holding the version number and the description of each migration it commits, and reads the highest version back at the next start-up to decide which migrations are still pending.

Adding New Migrations

To add a new migration, follow the steps below.

Edit collector/src/database/schema.go by appending a new Migration to the registerMigrations() method.
Set the version to the next free number, one higher than the last migration registered in the file.
Provide a clear, concise description of the migration; the description is stored in schema_version.
Implement the Up function, which receives the transaction the migration runs in.
Make the migration idempotent by using IF NOT EXISTS clauses where possible.
Describe every new object with COMMENT ON.

The examples below use 16 as the version number; replace it with the next free number at the time you write the migration.

Example: Adding a New Table

In the following example, the migration creates a new table:

sm.migrations = append(sm.migrations, Migration{
    Version:     16,
    Description: "Add probe_events table",
    Up: func(tx pgx.Tx) error {
        ctx := context.Background()

        _, err := tx.Exec(ctx, `
            CREATE TABLE IF NOT EXISTS probe_events (
                id BIGSERIAL PRIMARY KEY,
                connection_id INTEGER NOT NULL
                    REFERENCES connections(id)
                    ON DELETE CASCADE,
                collected_at TIMESTAMPTZ NOT NULL
                    DEFAULT CURRENT_TIMESTAMP,
                event_data JSONB NOT NULL
            );

            COMMENT ON TABLE probe_events IS
                'Discrete events reported by a probe.';
        `)
        if err != nil {
            return fmt.Errorf(
                "failed to create probe_events: %w", err,
            )
        }

        return nil
    },
})

Example: Adding an Index

In the following example, the migration creates an index on the collected_at column:

sm.migrations = append(sm.migrations, Migration{
    Version:     16,
    Description: "Add index on probe_events.collected_at",
    Up: func(tx pgx.Tx) error {
        ctx := context.Background()

        _, err := tx.Exec(ctx, `
            CREATE INDEX IF NOT EXISTS
                idx_probe_events_collected_at
                ON probe_events(collected_at DESC);

            COMMENT ON INDEX idx_probe_events_collected_at IS
                'Supports queries for the most recent events.';
        `)
        if err != nil {
            return fmt.Errorf(
                "failed to create index: %w", err,
            )
        }

        return nil
    },
})

PostgreSQL cannot build an index with CREATE INDEX CONCURRENTLY inside a transaction, and the SchemaManager wraps every migration in one, so an index added by a migration is built with a plain CREATE INDEX that blocks writes to the table whilst it runs.

Example: Adding a Constraint

PostgreSQL has no ADD CONSTRAINT IF NOT EXISTS, so schema.go provides the addConstraintIfMissing helper, which checks pg_constraint first and leaves an existing constraint untouched. In the following example, the migration adds a foreign key:

if err := addConstraintIfMissing(ctx, tx,
    "probe_events",
    "fk_probe_events_connection_id",
    "FOREIGN KEY (connection_id) "+
        "REFERENCES connections(id) ON DELETE CASCADE",
); err != nil {
    return err
}

Example: Modifying an Existing Column

In the following example, the migration adds a new column to an existing table:

sm.migrations = append(sm.migrations, Migration{
    Version:     16,
    Description: "Add priority column to probe_configs",
    Up: func(tx pgx.Tx) error {
        ctx := context.Background()

        _, err := tx.Exec(ctx, `
            ALTER TABLE probe_configs
                ADD COLUMN IF NOT EXISTS priority INTEGER
                NOT NULL DEFAULT 5;

            COMMENT ON COLUMN probe_configs.priority IS
                'Relative scheduling priority, 1 to 10.';
        `)
        if err != nil {
            return fmt.Errorf(
                "failed to add priority column: %w", err,
            )
        }

        return nil
    },
})

Best Practices

Follow these best practices when designing migrations and schema changes.

Migration Design

The following guidelines apply to migration design:

Include one logical change per migration; each migration should represent a single logical schema change.
Never modify applied migrations; create a new migration instead.
Make migrations idempotent; use IF NOT EXISTS, IF EXISTS, and existence checks.
Use transactions; the SchemaManager wraps each migration in a transaction.
Test migrations thoroughly on a development database before deploying.

Schema Design

The following guidelines apply to schema design:

Use constraints by defining CHECK, NOT NULL, UNIQUE, and FOREIGN KEY to enforce data integrity.
Create indexes strategically for foreign key columns, WHERE clause columns, ORDER BY columns, and JOIN conditions.
Use appropriate data types such as SERIAL for auto-incrementing IDs, TIMESTAMPTZ for timestamps, and TEXT for unlimited-length strings.
Include created_at and updated_at audit columns to track record modifications.
Plan for partitioning early for large tables such as metrics tables.

Testing

This section covers testing schema migrations.

Running Schema Tests

In the following example, the make command runs all tests:

make test

In the following example, the go test command runs only the migration tests:

go test -v -run TestMigrate

Test Environment

The schema tests need a PostgreSQL server they can create databases on. They do not run against the database named in the connection string; instead, TestMain calls setupTestDatabase(), which connects as an administrator, creates a database named ai_workbench_test_<YYYYMMDD_HHMMSS>_<microseconds>, and points the tests that connect through getTestConnection() at that generated database. One test, TestMigration_PgSettings, does not use the generated database; the Testing the pg_settings Migration section describes the exception.

Selecting a Server

The tests that use the generated database read TEST_AI_WORKBENCH_SERVER first and fall back to TEST_DB_CONN, which remains supported for older checkouts. Set the preferred variable to either a connection URL or a libpq key-value string, because replaceDatabase() accepts both forms. In the following example, the variable points the tests at the local PostgreSQL server:

export TEST_AI_WORKBENCH_SERVER=postgresql://[email protected]:5432/ai_workbench

Apart from TestMigration_PgSettings, the tests ignore the database name in the value, whether given as the URL path or as a dbname= field; the helpers rewrite the name to postgres for the administrative connection and to the generated name for the tests. The role in the connection string therefore needs the CREATEDB privilege, or superuser rights.

If neither variable is set, the helpers fall back to host=localhost port=5432 user=postgres sslmode=disable rather than skipping the tests. Always set TEST_AI_WORKBENCH_SERVER explicitly, so that the run resolves to the loopback server you intend.

Testing the pg_settings Migration

TestMigration_PgSettings in schema_pg_settings_test.go connects directly to the database named in TEST_AI_WORKBENCH_SERVER, rather than to the generated database. The test ignores TEST_DB_CONN and skips itself when TEST_AI_WORKBENCH_SERVER is unset. Before running the migrations, the test drops the metrics schema and the schema_version, probe_configs and connections tables in that database, so the value must name a disposable database on a loopback server, such as the empty ai_workbench database in the earlier example. Never point TEST_AI_WORKBENCH_SERVER at a database whose contents you need to keep, because this test destroys its schema on every run.

Keeping or Skipping Databases

Two further variables control what the run does with the database. In the following example, TEST_AI_WORKBENCH_KEEP_DB keeps the generated database after the tests finish, which helps when a migration test fails and the resulting schema needs inspection:

export TEST_AI_WORKBENCH_KEEP_DB=1

TEST_AI_WORKBENCH_KEEP_DB also accepts the value true. In the following example, SKIP_DB_TESTS skips the database tests altogether:

export SKIP_DB_TESTS=1

Leftover ai_workbench_test_* databases on a server come from runs that did not drop theirs, which happens whenever TEST_AI_WORKBENCH_KEEP_DB is set, whenever the run is interrupted before it reaches teardown, and whenever teardown itself fails, either because it cannot open its administrative connection or because DROP DATABASE returns an error. A complete run also leaves databases behind, because the three tests in maintenance_test.go each call setupTestDatabase() again; each call replaces the generated database for the rest of the run, and teardown drops only the last one. Leftover databases are safe to drop once you have confirmed that no test run or other client is still connected to them.

Confirming the Tests Ran

A passing run is not by itself evidence that the tests ran. When setupTestDatabase() fails for any reason, TestMain prints Skipping database tests and calls os.Exit(0), so the package reports success having run nothing; a wrong connection string produces a pass rather than an error. Run the tests with go test -v and check the output for that message, and for the individual test results, whenever a connection setting changes.

Writing Migration Tests

When adding a new migration, add corresponding tests that verify the following:

The migration applies successfully without errors.
Running the migration twice does not cause errors.
Constraints work as expected.
Indexes are created correctly.

In the following example, the test verifies that a migration creates a table:

func TestProbeEventsTable(t *testing.T) {
    pool, conn := getTestConnection(t)
    defer pool.Close()
    defer conn.Release()

    cleanupTestSchema(t, pool)
    sm := NewSchemaManager()
    if err := sm.Migrate(conn); err != nil {
        t.Fatalf("Failed to migrate: %v", err)
    }

    var count int
    err := pool.QueryRow(context.Background(), `
        SELECT COUNT(*)
        FROM information_schema.tables
        WHERE table_name = 'probe_events'
    `).Scan(&count)
    if err != nil {
        t.Fatalf("Failed to check for table: %v", err)
    }
    if count != 1 {
        t.Fatal("probe_events table was not created")
    }

    cleanupTestSchema(t, pool)
}

Troubleshooting

This section covers common schema management issues.

Migration Fails to Apply

If a migration fails, follow these steps:

Check the error message for details about the failure.
Verify that the database connection is accessible.
Review the migration code for logic errors.
Check for manual schema changes that conflict with the migration.

Migration Applied but Schema Incorrect

If a migration was applied but the schema is incorrect, follow these steps:

Check the schema_version table to verify which migrations were applied.
Investigate whether the migration partially applied before failing.
Create a new fix-up migration to correct the schema.

Rolling Back Migrations

The current system does not support automatic rollback. To roll back manually, follow these steps:

Use SQL to undo the migration changes manually.
Remove the migration record from schema_version.
Consider creating a new forward migration that reverts the changes instead.

Security Considerations

Follow these security practices when writing migrations.

Secure Migration Practices

The following guidelines apply to secure migrations:

Validate inputs if migrations use any configuration values.
Use parameterized queries when migration logic includes dynamic values.
Run migrations with a database user that has only the necessary privileges.
Review all migrations for security implications before applying.

Data Protection

The following guidelines apply to data protection:

Always back up the database before applying migrations in production.
Test migrations on a copy of production data before applying to production.
Ensure migrations do not inadvertently expose sensitive data.