Database functions¶

def add_data_to_db(sample_id: str, file_stem: str, start_uts: float, end_uts: float, job_id: str | None = None) -> str:
    """Register a time-series data file in the database.

    Warning: does not actually move a file. This function just updates the jobs and dataframes tables.

    Args:
        sample_id: Sample ID that the data is associated with
        file_stem: Filename of the file uploaded without snapshot. or extension
        start_uts: Data start unix time stamp
        end_uts: Data end unix time stamp
        job_id: Job ID that the data is associated with

    Returns:
        str: Job ID

    """
    data_start = datetime.fromtimestamp(start_uts, tz=timezone.utc).isoformat()
    data_end = datetime.fromtimestamp(end_uts, tz=timezone.utc).isoformat()
    if job_id:
        return _add_data_to_db_with_job(sample_id, file_stem, data_start, data_end, job_id)
    return _add_data_to_db_without_job(sample_id, file_stem, data_start, data_end)

def add_or_update_job(job_id: str, row: dict[str, str | float | None]) -> None:
    """Add or update job in database."""
    with engine.begin() as conn:
        conn.execute(
            insert(jobs_table)
            .values(stamp_sync({"Job ID": job_id, **row}, op="insert"))
            .on_conflict_do_update(
                index_elements=["Job ID"],
                set_=stamp_sync(row, op="update"),
            )
        )

def add_or_update_pipeline(pipeline: str, row: dict[str, str | float | None]) -> None:
    """Add or update pipeline in database. Remove job if ready == True."""
    # If ready is one, job gets removed
    if row.get("Ready") == 1:
        row["Job ID"] = None
        row["Job ID on server"] = None
    # If there is no Job ID, but there is a Job ID on the server, try to match it and add
    elif (
        isinstance(job_id_on_server := row.get("Job ID on server"), str)
        and isinstance(job_id := row.get("Server label"), str)
        and not row.get("Job ID")
    ):
        with suppress(ValueError):
            row["Job ID"] = get_job_id_from_server(job_id, job_id_on_server)
    # Insert or update the row
    uts = time()
    with engine.begin() as conn:
        conn.execute(
            insert(pipelines_table)
            .values(Pipeline=pipeline, **stamp_sync(row, uts=uts, op="insert"))
            .on_conflict_do_update(
                index_elements=["Pipeline"],
                set_=stamp_sync(row, uts=uts, op="update"),
            )
        )

def add_protocol_to_job(job_id: str, protocol: dict | str, capacity: float | None = None) -> None:
    """Attach a protocol to a job in the database."""
    if isinstance(protocol, dict):
        protocol = json.dumps(protocol)
    with engine.begin() as conn:
        conn.execute(
            update(jobs_table)
            .values(
                stamp_sync(
                    {
                        "Unicycler protocol": protocol,
                        "Capacity (mAh)": capacity,
                    }
                )
            )
            .where(jobs_table.c["Job ID"] == job_id)
        )

def add_samples_from_file(json_file: str | Path, overwrite: bool = False) -> None:
    """Add samples to database from a JSON file."""
    json_file = Path(json_file)
    _pre_check_sample_file(json_file)
    df = pd.read_json(json_file, orient="records")
    sample_df_to_db(df, overwrite)

def add_samples_from_object(samples: list[dict], overwrite: bool = False) -> None:
    """Add a samples to database from a list of dicts."""
    df = pd.DataFrame(samples)
    sample_df_to_db(df, overwrite)

def bulk_add_or_update_pipeline(rows: list[dict[str, str | float | None]]) -> None:
    """Add multiple rows to pipelines. Remove job if ready == True."""
    processed_rows = [
        {
            **row,
            **({"Job ID": None, "Job ID on server": None} if row.get("Ready") else {}),
        }
        for row in rows
    ]
    uts = time()
    with engine.begin() as conn:
        for row in processed_rows:
            conn.execute(
                insert(pipelines_table)
                .values(stamp_sync(row, uts=uts, op="insert"))
                .on_conflict_do_update(
                    index_elements=["Pipeline"],
                    set_=stamp_sync(row, uts=uts),
                )
            )

def check_job_running(job_id: str) -> bool:
    """Check if a job is currently on a pipeline."""
    with engine.connect() as conn:
        result = conn.execute(
            select(pipelines_table.c["Pipeline"]).where(pipelines_table.c["Job ID"] == job_id).limit(1)
        )
        return result.fetchone() is not None

def delete_samples(sample_ids: str | list) -> None:
    """Remove sample(s) from the database.

    Args:
        sample_ids: str or list
            The sample ID or list of sample IDs to remove from the database

    """
    if not isinstance(sample_ids, list):
        sample_ids = [sample_ids]
    with engine.begin() as conn:
        conn.execute(
            update(samples_table)
            .where(samples_table.c["Sample ID"].in_(sample_ids))
            .values(stamp_sync({}, op="delete"))
        )

def fill_pipelines_missing_job_ids() -> None:
    """Try to fill missing Job ID in pipelines if only Job ID on server is present."""
    job_id_subquery = (
        select(jobs_table.c["Job ID"])
        .where(jobs_table.c["Job ID on server"] == pipelines_table.c["Job ID on server"])
        .where(jobs_table.c["Server label"] == pipelines_table.c["Server label"])
        .scalar_subquery()
    )
    with engine.begin() as conn:
        conn.execute(
            update(pipelines_table)
            .where(pipelines_table.c["Job ID"].is_(None))
            .where(pipelines_table.c["Job ID on server"].isnot(None))
            .values(stamp_sync({"Job ID": job_id_subquery}, op="update"))
        )

def find_new_data(mode: str) -> list[str]:
    """Find jobs that have new data."""
    with engine.connect() as conn:
        if mode == "new_data":
            rows = conn.execute(
                select(
                    results_table.c["Sample ID"],
                    results_table.c["Last snapshot"],
                    results_table.c["Last analysis"],
                )
            ).fetchall()
            return [
                r[0] for r in rows if r[0] and (not r[1] or not r[2] or parse_datetime(r[1]) > parse_datetime(r[2]))
            ]
        if mode == "if_not_exists":
            rows = conn.execute(
                select(results_table.c["Sample ID"]).where(results_table.c["Last analysis"].is_(None))
            ).fetchall()
            return [r[0] for r in rows]
    return []

def get_all_run_ids() -> set[str]:
    """Get all valid run IDs."""
    with engine.connect() as conn:
        result = conn.execute(select(samples_table.c["Run ID"]).distinct()).fetchall()
    return {row[0] for row in result}

def get_all_sampleids() -> list[str]:
    """Get a list of all sample IDs in the database."""
    with engine.connect() as conn:
        result = conn.execute(select(samples_table.c["Sample ID"]).where(samples_table.c["sync_op"] != "delete"))
        return [row[0] for row in result.fetchall()]

def get_batch_details() -> dict[str, dict]:
    """Get all batch names, descriptions and samples from the database."""
    with engine.connect() as conn:
        result = conn.execute(
            select(batches_table.c.label, batches_table.c.description, batch_samples_table.c.sample_id)
            .join(batches_table, batch_samples_table.c.batch_id == batches_table.c.id)
            .order_by(batches_table.c.label)
        )
        batches: dict[str, dict] = {}
        for batch, description, sample in result.fetchall():
            if batch not in batches:
                batches[batch] = {"description": description, "samples": []}
            batches[batch]["samples"].append(sample)
        return dict(sorted(batches.items()))

def get_batches_from_sample(sample_id: str) -> list[str]:
    """Get the batch names that a sample belongs to."""
    with engine.connect() as conn:
        result = conn.execute(
            select(batches_table.c.label)
            .where(batch_samples_table.c.sample_id == sample_id)
            .join(batches_table, batch_samples_table.c.batch_id == batches_table.c.id)
        )
        return [r[0] for r in result.fetchall()]

def get_column_def(table: Table, column_names: list[str]) -> list[dict]:
    """Get AG grid definitions from a table and columns."""

    def map_type(col_type: Any) -> tuple[str, str]:  # noqa: ANN401
        """Get AG grid type and filter from sqlalchemy column type."""
        if isinstance(col_type, (Integer, Float, Numeric)):
            return "number", "agNumberColumnFilter"
        if isinstance(col_type, Boolean):
            return "boolean", "agTextColumnFilter"
        if isinstance(col_type, DateTime):
            return "date", "agDateColumnFilter"
        return "text", "agTextColumnFilter"

    columns = [table.columns[c] for c in column_names]
    col_types = [map_type(c.type) for c in columns]
    return [
        {
            "field": col.name,
            "cellDataType": col_type[0],
            "tooltipField": col.name,
            "filter": col_type[1],
            # Custom sorting for pipelines
            **(
                {"comparator": {"function": "pipelineComparatorCustom"}, "sort": "asc"}
                if col.name == "Pipeline"
                else {}
            ),
        }
        for col, col_type in zip(columns, col_types, strict=True)
    ]

def get_database(columns: dict[str, list] | None = None) -> dict[str, Any]:
    """Get all data from the database.

    Formatted for passing to Dash AG Grid rowData.
    """
    if columns is None:
        columns = {
            "samples": list(samples_table.columns.keys()),
            "pipelines": list(pipelines_table.columns.keys()),
            "jobs": list(jobs_table.columns.keys()),
            "results": list(results_table.columns.keys()),
        }
    with engine.connect() as conn:
        results = {
            "samples": conn.execute(
                select(*[samples_table.c[col] for col in columns["samples"]])
                .where(samples_table.c["sync_op"] != "delete")
                .order_by(samples_table.c["Sample ID"])
            ),
            "pipelines": conn.execute(  # Uses custom sort
                select(*[pipelines_table.c[col] for col in columns["pipelines"]]).where(
                    pipelines_table.c["sync_op"] != "delete"
                )
            ),
            "jobs": conn.execute(
                select(*[jobs_table.c[col] for col in columns["jobs"]])
                .where(jobs_table.c["sync_op"] != "delete")
                .order_by(jobs_table.c["Sample ID"])
            ),
            "results": conn.execute(
                select(*[results_table.c[col] for col in columns["results"]])
                .where(results_table.c["sync_op"] != "delete")
                .order_by(results_table.c["Sample ID"])
            ),
        }
        return {k: {"add": [dict(m) for m in v.mappings().all()]} for k, v in results.items()}

def get_database_updates(last_sync: float = 0, columns: dict[str, list] | None = None) -> dict[str, Any]:
    """Get all new data from the database.

    Formatted for viewing in Dash AG Grid.
    """
    if columns is None:
        columns = {
            "samples": list(samples_table.columns.keys()),
            "pipelines": list(pipelines_table.columns.keys()),
            "jobs": list(jobs_table.columns.keys()),
            "results": list(results_table.columns.keys()),
        }
    else:
        # Must collect these columns for basic functionality, even if they are not displayed
        pipelines_required = {"Pipeline", "Sample ID", "Job ID", "Server label"}
        columns["pipelines"] = list(set(columns["pipelines"]) | pipelines_required)
        if "Server label" not in columns["jobs"]:
            columns["jobs"].append("Server label")
    with engine.connect() as conn:
        results = {
            "samples": conn.execute(
                select(*[samples_table.c[col] for col in columns["samples"]])
                .where(samples_table.c["sync_modified"] > last_sync)
                .where(samples_table.c["sync_op"] != "delete")
                .order_by(samples_table.c["Sample ID"])
            ),
            "pipelines": conn.execute(
                select(*[pipelines_table.c[col] for col in columns["pipelines"]])
                .where(pipelines_table.c["sync_modified"] > last_sync)
                .where(pipelines_table.c["sync_op"] != "delete")
            ),
            "jobs": conn.execute(
                select(*[jobs_table.c[col] for col in columns["jobs"]])
                .where(jobs_table.c["sync_modified"] > last_sync)
                .where(jobs_table.c["sync_op"] != "delete")
                .order_by(jobs_table.c["Sample ID"])
            ),
            "results": conn.execute(
                select(*[results_table.c[col] for col in columns["results"]])
                .where(results_table.c["sync_modified"] > last_sync)
                .where(results_table.c["sync_op"] != "delete")
                .order_by(results_table.c["Sample ID"])
            ),
            "del_samples": conn.execute(
                select(samples_table.c["Sample ID"])
                .where(samples_table.c["sync_modified"] > last_sync)
                .where(samples_table.c["sync_op"] == "delete")
            ),
            "del_pipelines": conn.execute(
                select(pipelines_table.c["Pipeline"])
                .where(pipelines_table.c["sync_modified"] > last_sync)
                .where(pipelines_table.c["sync_op"] == "delete")
            ),
            "del_jobs": conn.execute(
                select(jobs_table.c["Job ID"])
                .where(jobs_table.c["sync_modified"] > last_sync)
                .where(jobs_table.c["sync_op"] == "delete")
            ),
            "del_results": conn.execute(
                select(results_table.c["Sample ID"])
                .where(results_table.c["sync_modified"] > last_sync)
                .where(results_table.c["sync_op"] == "delete")
            ),
        }
    return {
        table: {
            "upsert": [dict(m) for m in results[table].mappings().all()],
            "remove": [dict(m) for m in results[f"del_{table}"].mappings().all()],
        }
        for table in ["samples", "pipelines", "jobs", "results"]
    }

def get_db_last_update() -> float:
    """Get the last update time of the database.

    Returns 0.0 if never updated.
    """
    with engine.connect() as conn:
        return conn.execute(select(func.max(pipelines_table.c["sync_modified"]))).scalar() or 0.0

def get_job_data(job_id: str) -> dict:
    """Get all data about a job from the database."""
    with engine.connect() as conn:
        result = conn.execute(select(*job_cols).where(jobs_table.c["Job ID"] == job_id)).mappings().fetchone()
    if not result:
        msg = f"Job ID '{job_id}' not found in the database"
        raise ValueError(msg)
    job_data = dict(result)
    # Convert json strings to python objects
    payload = job_data.get("Payload")
    if payload and isinstance(payload, str) and payload.startswith(("[", "{")):
        job_data["Payload"] = json.loads(payload)
    unicycler = job_data.get("Unicycler protocol")
    if unicycler and isinstance(unicycler, str) and unicycler.startswith("{"):
        job_data["Unicycler protocol"] = json.loads(unicycler)

    return job_data

def get_job_from_pipeline(pipeline: str) -> str | None:
    """Get Job ID from a pipeline."""
    with engine.connect() as conn:
        return conn.execute(
            select(pipelines_table.c["Job ID"]).where(pipelines_table.c["Pipeline"] == pipeline)
        ).scalar()

def get_job_id_from_server(server_label: str, job_id_on_server: str) -> str:
    """Get the job ID from server label and job ID on server."""
    with engine.connect() as conn:
        result = conn.execute(
            select(jobs_table.c["Job ID"])
            .where(jobs_table.c["Job ID on server"] == job_id_on_server)
            .where(jobs_table.c["Server label"] == server_label)
        ).fetchone()
    if result:
        return result[0]
    msg = f"No Job ID found for server {server_label}: {job_id_on_server}"
    raise ValueError(msg)

def get_jobs_from_sample(sample_id: str) -> list[str]:
    """List all Job IDs associated with a sample."""
    with engine.connect() as conn:
        result = conn.execute(select(jobs_table.c["Job ID"]).where(jobs_table.c["Sample ID"] == sample_id)).fetchall()
    return [r[0] for r in result]

def get_last_harvest(server: dict, folder: str) -> float:
    """Get unix time stamp of last harvest."""
    with engine.connect() as conn:
        result = conn.execute(
            select(harvester_table.c["Last snapshot"])
            .where(harvester_table.c["Server label"] == server["label"])
            .where(harvester_table.c["Server hostname"] == server["hostname"])
            .where(harvester_table.c["Folder"] == folder)
        ).fetchone()
    if result:
        return parse_datetime(result[0]).timestamp()
    return 0.0

def get_neware_pipelines() -> tuple[list[str], list[str]]:
    """Get only running Neware pipelines."""
    with engine.connect() as conn:
        rows = conn.execute(
            select(pipelines_table.c["Pipeline"], pipelines_table.c["Server label"])
            .where(pipelines_table.c["Sample ID"].isnot(None))
            .where(pipelines_table.c["Ready"].is_(False))
            .where(pipelines_table.c["Server type"] == "neware")
        ).all()
    pipelines = [row[0] for row in rows]
    server_labels = [row[1] for row in rows]
    return pipelines, server_labels

def get_one_batch(batch_name: str) -> dict:
    """Get details of a batch from the batch name."""
    with engine.connect() as conn:
        result = conn.execute(
            select(batches_table.c.id, batches_table.c.description).where(batches_table.c.label == batch_name)
        )
        res = result.fetchone()
        if not res:
            msg = "Batch not in database"
            raise ValueError(msg)
        batch_id, description = res
        result = conn.execute(
            select(batch_samples_table.c.sample_id)
            .where(batch_samples_table.c.batch_id == batch_id)
            .order_by(batch_samples_table.c.sample_id)
        )
        samples = [s[0] for s in result.fetchall()]
    return {"name": batch_name, "description": description, "samples": samples}

def get_or_create_job_id_from_server(server_label: str, job_id_on_server: str) -> str:
    """Get the job ID from server label and job ID on server, create new Job ID if it doesn't exist."""
    try:
        job_id = get_job_id_from_server(server_label, job_id_on_server)
    except ValueError:
        job_id = str(uuid.uuid4())
        with engine.begin() as conn:
            conn.execute(
                insert(jobs_table).values(
                    stamp_sync(
                        {
                            "Job ID": job_id,
                            "Job ID on server": job_id_on_server,
                            "Server label": server_label,
                        },
                        op="insert",
                    )
                )
            )
    return job_id

def get_pipeline(pipeline: str) -> dict | None:
    """Get pipeline details."""
    with engine.connect() as conn:
        result = (
            conn.execute(select(*pipeline_cols).where(pipelines_table.c["Pipeline"] == pipeline)).mappings().first()
        )
    return dict(result) if result else None

def get_pipeline_from_sample(sample_id: str) -> dict | None:
    """Get pipeline from a Sample ID."""
    with engine.connect() as conn:
        result = (
            conn.execute(select(*pipeline_cols).where(pipelines_table.c["Sample ID"] == sample_id)).mappings().first()
        )
    return dict(result) if result else None

def get_results_from_sample(sample_id: str) -> dict | None:
    """Get results summary from Sample ID."""
    with engine.connect() as conn:
        result = conn.execute(select(*result_cols).where(results_table.c["Sample ID"] == sample_id)).mappings().first()
    return dict(result) if result else None

def get_running_job(sample_id: str) -> dict[str, str | None]:
    """Get pipeline, job ID, and status of a job if a sample is running."""
    with engine.connect() as conn:
        result = (
            conn.execute(
                select(
                    pipelines_table.c["Pipeline"],
                    pipelines_table.c["Job ID"],
                    jobs_table.c["Status"],
                )
                .outerjoin(jobs_table, pipelines_table.c["Job ID"] == jobs_table.c["Job ID"])
                .where(pipelines_table.c["Sample ID"] == sample_id)
            )
            .mappings()
            .fetchone()
        )

    if result:
        return dict(result)
    return {"Pipeline": None, "Job ID": None, "Status": None}

def get_sample_data(sample_id: str) -> dict:
    """Get all data about a sample from the database."""
    with engine.connect() as conn:
        result = (
            conn.execute(
                select(*sample_cols)
                .where(samples_table.c["Sample ID"] == sample_id)
                .where(samples_table.c["sync_op"] != "delete")
            )
            .mappings()
            .fetchone()
        )
        if not result:
            msg = f"Sample ID '{sample_id}' not found in the database"
            raise ValueError(msg)
        sample_data = dict(result)
        # Convert json strings to python objects
        history = sample_data.get("Assembly history")
        if history and isinstance(history, str):
            sample_data["Assembly history"] = json.loads(history)
    return sample_data

def get_sample_from_pipeline(pipeline: str) -> str | None:
    """Get Sample ID from a pipeline."""
    with engine.connect() as conn:
        return conn.execute(
            select(pipelines_table.c["Sample ID"]).where(pipelines_table.c["Pipeline"] == pipeline)
        ).scalar()

def get_unicycler_protocols(sample_id: str) -> list[dict]:
    """Return a list of unicycler protocols associated with the sample."""
    with engine.connect() as conn:
        j = jobs_table.c
        d = dataframes_table.c

        sort_timestamp = func.coalesce(
            j["Submitted"],
            select(func.min(d["Data start"])).where(d["Job ID"] == j["Job ID"]).scalar_subquery(),
            literal(datetime(9999, 12, 31), type_=j["Submitted"].type),  # noqa: DTZ001
        ).label("sort_timestamp")

        result = conn.execute(
            select(j["Job ID"], j["Unicycler protocol"], j["Capacity (mAh)"], sort_timestamp)
            .where(j["Sample ID"] == sample_id)
            .where(j["Unicycler protocol"].isnot(None))
            .order_by(sort_timestamp)
        )
    return [dict(row) for row in result.mappings().all()]

def is_sample(sample_id: str) -> bool:
    """Check if `sample_id` exists in the database."""
    with engine.connect() as conn:
        return bool(conn.execute(select(exists().where(samples_table.c["Sample ID"] == sample_id))).scalar())

def patch_database(engine: Engine) -> None:
    """Add missing columns to database, in case users are coming from an older version."""
    if _db_schema_needs_update(engine):
        try:
            _update_db_schema(engine)
        except ProgrammingError as e:
            msg = (
                "Database schema needs updating. "
                "Failed to update. An admin must run 'aurora-app' or 'aurora-setup update' first."
            )
            raise PermissionError(msg) from e

def remove_batch(batch_name: str) -> None:
    """Remove a batch from the database."""
    with engine.begin() as conn:
        batch_id = conn.execute(select(batches_table.c.id).where(batches_table.c.label == batch_name)).fetchone()[0]
        conn.execute(delete(batches_table).where(batches_table.c.label == batch_name))
        conn.execute(delete(batch_samples_table).where(batch_samples_table.c.batch_id == batch_id))

def sample_df_to_db(df: pd.DataFrame, overwrite: bool = False) -> None:
    """Add samples to database from a pandas dataframe."""
    sample_ids = df["Sample ID"].tolist()
    if any(not isinstance(sample_id, str) for sample_id in sample_ids):
        msg = "File contains non-string 'Sample ID' keys"
        raise TypeError(msg)
    for sample_id in sample_ids:
        check_illegal_text(sample_id)
    if len(sample_ids) != len(set(sample_ids)):
        msg = "File contains duplicate 'Sample ID' keys"
        raise ValueError(msg)

    # Check if any sample already exists
    existing_sample_ids = get_all_sampleids()
    if not overwrite and any(sample_id in existing_sample_ids for sample_id in sample_ids):
        msg = "Sample IDs already exist in the database"
        raise ValueError(msg)

    # Recalculate some values
    df = _recalculate_sample_data(df)

    # Insert into database
    valid_columns = {c.name for c in samples_table.columns}
    df_columns = set(df.columns)
    missing_in_db = df_columns - valid_columns
    if missing_in_db:
        logger.warning(
            "Adding samples to database: after automatic calculations, "
            "column(s) %s do not exist in the database, skipping",
            ", ".join("'" + col + "'" for col in missing_in_db),
        )
        df = df.drop(columns=missing_in_db)

    with engine.begin() as conn:
        for _, raw_row in df.iterrows():
            # Remove empty columns from the row
            row = raw_row.dropna()

            # Insert or update the row
            row_dict = row.to_dict()
            conn.execute(
                insert(samples_table)
                .values(stamp_sync(row_dict, op="insert"))
                .on_conflict_do_update(
                    index_elements=["Sample ID"],
                    set_=stamp_sync(row_dict),
                )
            )

def save_or_overwrite_batch(batch_name: str, batch_description: str, sample_ids: list, overwrite: bool = False) -> None:
    """Save a batch to the database, overwriting it if the name already exists."""
    with engine.begin() as conn:
        result = conn.execute(select(batches_table.c.id).where(batches_table.c.label == batch_name)).fetchone()

        if result:
            if not overwrite:
                msg = f"Batch {batch_name} already exists. Set overwrite=True to overwrite."
                raise ValueError(msg)
            batch_id = result[0]
            conn.execute(
                update(batches_table).where(batches_table.c.id == batch_id).values(description=batch_description)
            )
            conn.execute(delete(batch_samples_table).where(batch_samples_table.c.batch_id == batch_id))
        else:
            batch_id = conn.execute(
                insert(batches_table).values(label=batch_name, description=batch_description)
            ).inserted_primary_key[0]

        conn.execute(
            insert(batch_samples_table),
            [{"batch_id": batch_id, "sample_id": sample_id} for sample_id in sample_ids],
        )