Files
agentic-mobile-control/storage/device_config.py
T

155 lines
5.2 KiB
Python

from __future__ import annotations
import json
import sqlite3
from pathlib import Path
from typing import Any
DEFAULT_MAX_STEPS = 20
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"],
}