#!/usr/bin/env python3
"""Limited interace to MLflow's SQLite database."""

import pathlib
import sqlite3


class MLflowSqlite:
    """
    Limited interface to MLflow's SQLite database.
    """

    DEFAULT_PATH = "mlflow.db"

    def __init__(self, path=DEFAULT_PATH):
        self.path = pathlib.Path(path).resolve()
        self._conn = None

    @property
    def conn(self):
        """
        The connection.
        """
        return sqlite3.connect(self.path)

    def update_artifact_uris(self, artifact_map):
        """
        Update artifact URIs in the database.

        Args:
            artifact_map:
                A dict mapping run IDs to new artifact URIs.
        """
        args = [(v, k) for (k, v) in artifact_map.items()]
        with self.conn as conn:
            conn.executemany(
                "UPDATE runs SET artifact_uri = ? WHERE run_uuid = ?", args
            )
