agentic-mobile-control/storage/device_config.py

from __future__ import annotations

import json
import sqlite3
from pathlib import Path
from typing import Any


DEFAULT_MAX_STEPS = 999999


class DeviceConfigStore:
    def __init__(
        self,
        db_path: str | Path = "tasks/device_config.sqlite3",
        *,
        default_max_steps: int = DEFAULT_MAX_STEPS,
    ) -> None:
        self.db_path = Path(db_path)
        self.default_max_steps = default_max_steps
        self.db_path.parent.mkdir(parents=True, exist_ok=True)
        self._ensure_schema()

    def add(
        self,
        *,
        device_id: str,
        name: str | None = None,
        driver_type: str,
        connection_info: dict[str, Any],
        cloud_device_id: str | None = None,
    ) -> None:
        with self._connect() as connection:
            connection.execute(
                """
                insert into device_configs (
                    device_id, name, driver_type, connection_info, cloud_device_id
                ) values (?, ?, ?, ?, ?)
                on conflict(device_id) do update set
                    name = excluded.name,
                    driver_type = excluded.driver_type,
                    connection_info = excluded.connection_info,
                    cloud_device_id = coalesce(
                        excluded.cloud_device_id,
                        device_configs.cloud_device_id
                    )
                """,
                (
                    device_id,
                    name,
                    driver_type,
                    json.dumps(connection_info, ensure_ascii=False),
                    cloud_device_id,
                ),
            )

    def set_cloud_device_id(self, device_id: str, cloud_device_id: str) -> bool:
        if not cloud_device_id:
            raise ValueError("cloud_device_id must not be empty")
        with self._connect() as connection:
            cursor = connection.execute(
                "update device_configs set cloud_device_id = ? where device_id = ?",
                (cloud_device_id, device_id),
            )
        return cursor.rowcount > 0

    def remove(self, device_id: str) -> None:
        with self._connect() as connection:
            connection.execute(
                "delete from device_configs where device_id = ?",
                (device_id,),
            )

    def get(self, device_id: str) -> dict[str, Any] | None:
        with self._connect() as connection:
            row = connection.execute(
                "select * from device_configs where device_id = ?",
                (device_id,),
            ).fetchone()
        return self._row_to_config(row) if row else None

    def list(self) -> list[dict[str, Any]]:
        with self._connect() as connection:
            rows = connection.execute(
                "select * from device_configs order by device_id"
            ).fetchall()
        return [self._row_to_config(row) for row in rows]

    def get_setting(self, key: str) -> str | None:
        with self._connect() as connection:
            row = connection.execute(
                "select value from settings where key = ?",
                (key,),
            ).fetchone()
        return str(row["value"]) if row else None

    def set_setting(self, key: str, value: object) -> None:
        with self._connect() as connection:
            connection.execute(
                """
                insert into settings (key, value)
                values (?, ?)
                on conflict(key) do update set value = excluded.value
                """,
                (key, str(value)),
            )

    def _ensure_schema(self) -> None:
        with self._connect() as connection:
            connection.execute(
                """
                create table if not exists device_configs (
                    device_id text primary key,
                    name text,
                    driver_type text not null,
                    connection_info text not null,
                    cloud_device_id text
                )
                """
            )
            columns = {
                row["name"]
                for row in connection.execute("pragma table_info(device_configs)")
            }
            if "cloud_device_id" not in columns:
                connection.execute(
                    "alter table device_configs add column cloud_device_id text"
                )
            connection.execute(
                """
                create table if not exists settings (
                    key text primary key,
                    value text not null
                )
                """
            )
            connection.execute(
                "insert or ignore into settings (key, value) values (?, ?)",
                ("max_steps", str(self.default_max_steps)),
            )

    def _connect(self) -> sqlite3.Connection:
        connection = sqlite3.connect(self.db_path)
        connection.row_factory = sqlite3.Row
        return connection

    def _row_to_config(self, row: sqlite3.Row) -> dict[str, Any]:
        return {
            "device_id": row["device_id"],
            "name": row["name"],
            "driver_type": row["driver_type"],
            "connection_info": json.loads(row["connection_info"]),
            "cloud_device_id": row["cloud_device_id"],
        }