v3/server/bin/create-election.py

#!/usr/bin/env -S uv run --script

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

"""
Script to create an election from a YAML definition file.
Reads election metadata, issues, and voter records, then populates the database.
"""

import argparse
import datetime
import pathlib
import logging
import yaml

import steve.election
import steve.persondb
import steve.vtypes.stv

_LOGGER = logging.getLogger(__name__)

THIS_DIR = pathlib.Path(__file__).resolve().parent
DB_FNAME = THIS_DIR.parent / 'steve.db'

# Supported vote types
VALID_VTYPES = {'yna', 'stv'}


def parse_datetime(dt):
    """Parse ISO datetime string to Unix timestamp, or return if already int."""
    if not dt:
        return None
    if isinstance(dt, int):
        return dt
    if isinstance(dt, str):
        dt_obj = datetime.datetime.fromisoformat(dt)
        return int(dt_obj.timestamp())
    raise ValueError(f'Invalid datetime format: {dt}')


def validate_issue(issue):
    """Validate an issue dict from YAML."""
    if 'vtype' not in issue or issue['vtype'] not in VALID_VTYPES:
        raise ValueError(f'Invalid vtype: {issue.get("vtype")}')
    if issue['vtype'] == 'stv':
        kv = issue.get('kv', {})
        if not all(k in kv for k in ['version', 'labelmap', 'seats']):
            raise ValueError(
                'STV issue missing required kv fields: version, labelmap, seats'
            )
        if not isinstance(kv['seats'], int) or kv['seats'] <= 0:
            raise ValueError('STV seats must be a positive integer')
        steve.vtypes.stv.get_candidates(kv)
    return issue


def main(yaml_file):
    with open(yaml_file, 'r') as f:
        data = yaml.safe_load(f)

    # Extract election data
    election_data = data.get('election', {})
    title = election_data.get('title')
    owner_pid = election_data.get('owner_pid')
    authz = election_data.get('authz')
    open_at = parse_datetime(election_data.get('open_at'))
    close_at = parse_datetime(election_data.get('close_at'))

    if not title or not owner_pid:
        raise ValueError('Election must have title and owner_pid')

    # Extract issues
    issues = data.get('issues', [])
    for issue in issues:
        validate_issue(issue)

    # Extract record: list of pid values for eligible voters
    record = data.get('record', [])
    if not isinstance(record, list):
        raise ValueError('record must be a list of pid values')

    # TODO: Re-enable transactions for safety once database setup allows.
    # pdb = steve.persondb.PersonDB(DB_FNAME)
    # pdb.db.conn.execute('BEGIN TRANSACTION')

    try:
        # Create election
        election = steve.election.Election.create(
            DB_FNAME, title, owner_pid, authz, open_at, close_at
        )
        _LOGGER.info(
            f'Created election[E:{election.eid}]: "{title}" by owner "{owner_pid}"'
        )

        # Add issues
        for issue_data in issues:
            iid = election.add_issue(
                issue_data['title'],
                issue_data.get('description'),
                issue_data['vtype'],
                issue_data.get('kv') if issue_data['vtype'] == 'stv' else None,
            )
            _LOGGER.info(f'Added issue[I:{iid}] to election[E:{election.eid}]')

        # Open a PersonDB using the existing DB from the Election
        pdb = steve.persondb.PersonDB(election.db)

        # HACK: Opened PDB using existing DB; lacks PersonDB cursors.
        # q_person: SELECT * FROM person ORDER BY pid
        pdb.q_person = pdb.db.cursor_for('SELECT * FROM person ORDER BY pid')

        # Get all persons for validation
        all_persons = pdb.list_persons()
        all_pids = {person.pid for person in all_persons}

        # Temporary hack: Map old PIDs to newer equivalents for testing.
        # TODO: Remove once PID data is fully migrated.
        _REMAP = {
            'iroh': 'wells',
        }

        # Validate and add voters from record
        for pid in record:
            pid = _REMAP.get(pid, pid)  # Apply remapping if needed
            if pid not in all_pids:
                raise ValueError(f'PID {pid} from record not found in person database')
            election.add_voter(pid)
        _LOGGER.info(f'Added {len(record)} voters to election[E:{election.eid}]')

        # pdb.db.conn.execute('COMMIT')
        _LOGGER.info(f'Election[E:{election.eid}] fully created from {yaml_file}')

    except Exception as e:
        # pdb.db.conn.execute('ROLLBACK')
        _LOGGER.error(f'Failed to create election from {yaml_file}: {e}')
        raise


if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO)

    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument(
        'yaml_file',
        help='Path to the YAML file defining the election.',
    )
    args = parser.parse_args()
    main(args.yaml_file)