Example schema and usage for https://getvet.sh

get_vet_schema_usage.py

import gzip
import hashlib
import sqlite3
import urllib.request
from dataclasses import asdict, dataclass


@dataclass
class Queries:
    insert_program: str = (
        "INSERT OR IGNORE INTO Program (name, homepage, script_url) "
        "VALUES (:name, :homepage, :script_url) "
        "RETURNING id"
    )
    insert_program_version: str = (
        "INSERT OR IGNORE INTO ProgramVersion (program_id, program_version, script_hash) "
        "VALUES (:program_id, :program_version, :script_hash) "
        "RETURNING id"
    )
    insert_script: str = (
        "INSERT OR IGNORE INTO Script (program_version_id, body) "
        "VALUES (:program_version_id, :body) "
        "RETURNING id"
    )
    select_program: str = (
        "SELECT id FROM Program WHERE name = :name AND script_url = :script_url"
    )

    select_program_version_by_id_ver: str = "SELECT id FROM ProgramVersion WHERE program_id = :program_id AND program_version = :program_version"

    select_program_version_by_hash: str = (
        "SELECT id FROM ProgramVersion WHERE script_hash = :script_hash"
    )


@dataclass
class Script:
    name: str
    body: bytes | None = None
    homepage: str | None = None
    program_id: int | None = None
    program_version: str | None = None
    program_version_id: int | None = None
    script_hash: str | None = None
    script_url: str | None = None


def get_conn(db_name: str) -> sqlite3.Connection:
    conn = sqlite3.connect(db_name)
    conn.execute("PRAGMA foreign_keys = ON")

    return conn


def create_schema(conn: sqlite3.Connection) -> None:
    program_schema = (
        "CREATE TABLE IF NOT EXISTS Program (\n"
        "  id INTEGER NOT NULL PRIMARY KEY,\n"
        "  name TEXT NOT NULL,\n"
        "  homepage TEXT DEFAULT NULL,\n"
        "  script_url TEXT NOT NULL,\n"
        "  CONSTRAINT Program_script_url_unq\n"
        "    UNIQUE (name COLLATE NOCASE, script_url COLLATE NOCASE)\n"
        ") STRICT"
    )

    script_schema = (
        "CREATE TABLE IF NOT EXISTS Script (\n"
        "  id INTEGER NOT NULL PRIMARY KEY,\n"
        "  program_version_id INTEGER NOT NULL UNIQUE,\n"
        "  body BLOB NOT NULL,\n"
        "  CONSTRAINT ProgramVersion_Script_id_fkey\n"
        "    FOREIGN KEY (program_version_id)\n"
        "    REFERENCES ProgramVersion (id)\n"
        "      ON UPDATE RESTRICT\n"
        "      ON DELETE RESTRICT\n"
        "      DEFERRABLE INITIALLY DEFERRED\n"
        ") STRICT"
    )

    program_version_schema = (
        "CREATE TABLE IF NOT EXISTS ProgramVersion (\n"
        "  id INTEGER NOT NULL PRIMARY KEY,\n"
        "  program_id INTEGER NOT NULL,\n"
        "  program_version TEXT NOT NULL,\n"
        "  script_hash TEXT NOT NULL,\n"
        "  download_count INTEGER NOT NULL DEFAULT 0,\n"
        "  created_at TEXT NOT NULL DEFAULT (CURRENT_TIMESTAMP),\n"
        "  CONSTRAINT Program_ProgramVersion_fkey\n"
        "    FOREIGN KEY (program_id) REFERENCES Program (id)\n"
        "      ON UPDATE RESTRICT\n"
        "      ON DELETE CASCADE\n"
        "      DEFERRABLE INITIALLY DEFERRED"
        "  CONSTRAINT ProgramVersion_script_hash_unq\n"
        "    UNIQUE (script_hash COLLATE NOCASE)\n"
        ") STRICT"
    )

    with conn:
        for schema in (program_schema, script_schema, program_version_schema):
            conn.execute(schema)


def create_indices(conn: sqlite3.Connection) -> None:
    with conn:
        conn.execute(
            "CREATE INDEX IF NOT EXISTS ProgramVersion_id_version_idx ON "
            "ProgramVersion (program_id, program_version COLLATE NOCASE)"
        )


def hash_script(script_body: bytes) -> str:
    m = hashlib.sha256()
    m.update(script_body)

    return m.hexdigest()


def download_script(url: str) -> bytes | None:
    # if this fails, let it fail
    with urllib.request.urlopen(url) as resp:
        return resp.read()

    return None


def compress_script(script_body: bytes) -> bytes:
    return gzip.compress(script_body)


def insert_new_script(conn: sqlite3.Connection, script: Script) -> None:
    cur = conn.cursor()

    cur.execute(Queries.insert_program, asdict(script))
    if not (program_id := cur.fetchone()):
        print(f"DEBUG: UNIQUE tuple ({script.name}, {script.script_url}) exists")
        cur.execute(Queries.select_program, asdict(script))
        program_id = cur.fetchone()
    assert program_id
    script.program_id = program_id[0]

    cur.execute(Queries.insert_program_version, asdict(script))
    if not (program_version_id := cur.fetchone()):
        print(f"DEBUG: UNIQUE tuple ({script.script_hash}) exists")
        cur.execute(Queries.select_program_version_by_hash, asdict(script))
        program_version_id = cur.fetchone()
    assert program_version_id
    script.program_version_id = program_version_id[0]

    cur.execute(Queries.insert_script, asdict(script))
    if not cur.fetchone():
        print(f"DEBUG: UNIQUE tuple ({script.program_version_id}) exists")

    conn.commit()
    cur.close()

    return None


def make_script(conn: sqlite3.Connection) -> Script:
    script = Script(name=input("Program's name: "))
    script.program_version = input("Program's version: ")
    script.script_url = input("Program's installer script URL: ")
    script.homepage = input("Program's homepage: ") or None
    if (_script_body := download_script(script.script_url)) is None:
        raise ValueError(f"Script {script.name} had no body")
    script.body = compress_script(_script_body)
    script.script_hash = hash_script(_script_body)

    return script


if __name__ == "__main__":
    conn = get_conn("vet.db")
    create_schema(conn)
    create_indices(conn)
    script = make_script(conn)
    insert_new_script(conn, script)