lablink_allocator_service.generate_init_sql

lablink_allocator_service.generate_init_sql

Attributes

Functions:

build_init_sql()

Build the PostgreSQL initialization SQL script as a string.

Source code in packages/allocator/src/lablink_allocator_service/generate_init_sql.py

def build_init_sql() -> str:
    """Build the PostgreSQL initialization SQL script as a string."""
    config = get_config()
    DB_PASSWORD = config.db.password
    # MESSAGE_CHANNEL was the CRD notify channel; no longer needed.

    return f"""
ALTER SYSTEM SET listen_addresses = '*';

SET client_min_messages TO WARNING;
DROP USER IF EXISTS {DB_USER};
CREATE USER {DB_USER} WITH ENCRYPTED PASSWORD '{DB_PASSWORD}';
ALTER USER {DB_USER} WITH LOGIN;
CREATE DATABASE {DB_NAME} OWNER {DB_USER};
GRANT ALL PRIVILEGES ON DATABASE {DB_NAME} TO {DB_USER};

-- For /api/health/connections: without it, pg_stat_activity masks other roles'
-- `state` as NULL. Read-only access to the stats views, nothing else.
GRANT pg_read_all_stats TO {DB_USER};

\\c {DB_NAME};

SET ROLE {DB_USER};

CREATE TABLE IF NOT EXISTS scheduled_destructions (
    id SERIAL PRIMARY KEY,
    schedule_name VARCHAR(255) NOT NULL UNIQUE,
    destruction_time TIMESTAMP NOT NULL,
    recurrence_rule VARCHAR(255),
    created_by VARCHAR(255),
    status VARCHAR(50) NOT NULL DEFAULT 'scheduled',
    execution_count INTEGER DEFAULT 0,
    last_execution_time TIMESTAMP,
    last_execution_result TEXT,
    notification_enabled BOOLEAN DEFAULT TRUE,
    notification_hours_before INTEGER DEFAULT 1,
    created_at TIMESTAMP DEFAULT NOW(),
    updated_at TIMESTAMP DEFAULT NOW()
);

CREATE INDEX idx_destruction_time ON scheduled_destructions(destruction_time);
CREATE INDEX idx_status ON scheduled_destructions(status);

CREATE OR REPLACE FUNCTION update_scheduled_destructions_updated_at()
RETURNS TRIGGER AS $$
BEGIN
    NEW.updated_at = NOW();
    RETURN NEW;
END;
$$ LANGUAGE plpgsql;

CREATE TRIGGER scheduled_destructions_updated_at
    BEFORE UPDATE ON scheduled_destructions
    FOR EACH ROW
    EXECUTE FUNCTION update_scheduled_destructions_updated_at();

CREATE TABLE IF NOT EXISTS operations (
    id SERIAL PRIMARY KEY,
    op_type VARCHAR(16) NOT NULL,
    status VARCHAR(16) NOT NULL,
    params TEXT,
    created_by VARCHAR(255),
    created_at TIMESTAMP NOT NULL DEFAULT NOW(),
    started_at TIMESTAMP,
    finished_at TIMESTAMP,
    output TEXT,
    error TEXT,
    resources_total INTEGER,
    resources_completed INTEGER
);

CREATE INDEX idx_operations_created_at ON operations(created_at);

-- At most one queued/running operation at a time. This is the entire
-- concurrency guard: two near-simultaneous INSERTs under Flask's
-- threaded=True can't both succeed, so no separate locking code is needed.
CREATE UNIQUE INDEX IF NOT EXISTS operations_single_flight
    ON operations ((1)) WHERE status IN ('queued', 'running');

CREATE TABLE IF NOT EXISTS {VM_TABLE_NAME} (
    HostName VARCHAR(1024) PRIMARY KEY,
    UserEmail VARCHAR(1024),
    InUse BOOLEAN NOT NULL DEFAULT FALSE,
    Healthy VARCHAR(1024),
    Status   VARCHAR(1024),
    CloudInitLogs TEXT,
    DockerLogs TEXT,
    TofuApplyStartTime TIMESTAMP,
    TofuApplyEndTime TIMESTAMP,
    TofuApplyDurationSeconds FLOAT,
    CloudInitStartTime TIMESTAMP,
    CloudInitEndTime TIMESTAMP,
    CloudInitDurationSeconds FLOAT,
    ContainerStartTime TIMESTAMP,
    ContainerEndTime TIMESTAMP,
    ContainerStartupDurationSeconds FLOAT,
    TotalStartupDurationSeconds FLOAT,
    CreatedAt TIMESTAMP DEFAULT NOW(),
    last_seen_at TIMESTAMP,
    boot_id VARCHAR(64),
    disk_free_pct SMALLINT,
    SessionId UUID,
    BrowserToken TEXT,
    VncPassword TEXT,
    Upstream TEXT,
    SessionStartedAt TIMESTAMPTZ,
    AdminReservedAt TIMESTAMPTZ,
    machine_identity   TEXT,
    provider           TEXT NOT NULL DEFAULT 'aws',
    endpoint_url       TEXT,
    provider_metadata  JSONB NOT NULL DEFAULT '{{}}',
    client_secret_hash TEXT,
    browser_ws_url     TEXT,
    browser_credential TEXT,
    gpu_present        BOOLEAN,
    gpu_model          TEXT,
    last_release_time  TIMESTAMP,
    SessionMetricsStartedAt          TIMESTAMPTZ,
    SessionMetricsSealedAt           TIMESTAMPTZ,
    SessionMetricsLastReportedAt     TIMESTAMPTZ,
    SecondsInSubjectSoftware         INTEGER,
    SecondsInTerminal                INTEGER,
    SecondsInBrowser                 INTEGER,
    SecondsInOther                   INTEGER,
    GpuActiveSeconds                 INTEGER,
    GpuUtilPeak                      INTEGER,
    VramUsedPeakMb                   INTEGER,
    SecondsToFirstSleapLabel         INTEGER,
    SecondsToFirstSleapTrain         INTEGER,
    SecondsToFirstSleapTrack         INTEGER,
    MaxLabeledFrames                 INTEGER,
    TrainingEpochsCompleted          INTEGER,
    TrainingFinalLoss                DOUBLE PRECISION,
    SessionMetricsRaw                JSONB
);

CREATE UNIQUE INDEX {VM_TABLE_NAME}_browser_token_idx
    ON {VM_TABLE_NAME}(BrowserToken) WHERE BrowserToken IS NOT NULL;
CREATE UNIQUE INDEX {VM_TABLE_NAME}_session_id_idx
    ON {VM_TABLE_NAME}(SessionId) WHERE SessionId IS NOT NULL;
CREATE UNIQUE INDEX {VM_TABLE_NAME}_machine_identity_idx
    ON {VM_TABLE_NAME}(machine_identity) WHERE machine_identity IS NOT NULL;
CREATE INDEX {VM_TABLE_NAME}_provider_idx ON {VM_TABLE_NAME}(provider);
CREATE INDEX {VM_TABLE_NAME}_assignable_idx
    ON {VM_TABLE_NAME}(status, useremail, last_release_time)
    WHERE useremail IS NULL AND status = 'running' AND adminreservedat IS NULL;

CREATE TABLE IF NOT EXISTS settings (
    key TEXT PRIMARY KEY,
    value TEXT NOT NULL
);
"""

main()

Write init.sql to the path Postgres' bootstrap step reads.

Source code in packages/allocator/src/lablink_allocator_service/generate_init_sql.py

def main():
    """Write init.sql to the path Postgres' bootstrap step reads."""
    with open("/app/init.sql", "w") as f:
        f.write(build_init_sql())