1  """Functions for syncing data from an external database of manufacturers, filaments, materials, etc."""
  2  
  3  import datetime
  4  import logging
  5  import os
  6  from collections.abc import Iterator
  7  from enum import Enum
  8  from pathlib import Path
  9  from urllib.parse import urljoin
 10  
 11  import hishel
 12  from pydantic import BaseModel, Field, RootModel
 13  from scheduler.asyncio.scheduler import Scheduler
 14  
 15  from spoolman import filecache
 16  from spoolman.env import get_cache_dir
 17  
 18  logger = logging.getLogger(__name__)
 19  
 20  
 21  DEFAULT_EXTERNAL_DB_URL = "https://donkie.github.io/SpoolmanDB/"
 22  DEFAULT_SYNC_INTERVAL = 3600
 23  
 24  controller = hishel.Controller(allow_stale=True)
 25  try:
 26      cache_path = get_cache_dir() / "hishel"
 27      cache_storage = hishel.AsyncFileStorage(base_path=cache_path)
 28  except PermissionError:
 29      logger.warning(
 30          "Failed to setup disk-based cache due to permission error. Ensure the path %s is writable. "
 31          "Using in-memory cache instead as fallback.",
 32          str(cache_path.resolve()),
 33      )
 34      cache_storage = hishel.AsyncInMemoryStorage()
 35  
 36  
 37  class SpoolType(Enum):
 38      PLASTIC = "plastic"
 39      CARDBOARD = "cardboard"
 40      METAL = "metal"
 41  
 42  
 43  class Finish(Enum):
 44      MATTE = "matte"
 45      GLOSSY = "glossy"
 46  
 47  
 48  class MultiColorDirection(Enum):
 49      COAXIAL = "coaxial"
 50      LONGITUDINAL = "longitudinal"
 51  
 52  
 53  class Pattern(Enum):
 54      MARBLE = "marble"
 55      SPARKLE = "sparkle"
 56  
 57  
 58  class ExternalFilament(BaseModel):
 59      id: str = Field(description="A unique ID for this filament.", examples=["polymaker_pla_polysonicblack_1000_175"])
 60      manufacturer: str = Field(description="Filament manufacturer.", examples=["Polymaker"])
 61      name: str = Field(description="Filament name.", examples=["Polysonic\u2122 Black"])
 62      material: str = Field(description="Filament material.", examples=["PLA"])
 63      density: float = Field(description="Density in g/cm3.", examples=[1.23])
 64      weight: float = Field(description="Net weight of a single spool.", examples=[1000])
 65      spool_weight: float | None = Field(default=None, description="Weight of an empty spool.", examples=[140])
 66      spool_type: SpoolType | None = Field(None, description="Type of spool.", examples=[SpoolType.PLASTIC])
 67      diameter: float = Field(description="Filament in mm.", examples=[1.75])
 68      color_hex: str | None = Field(
 69          default=None,
 70          description="Filament color code in hex format, for single-color filaments.",
 71          examples=["2c3232"],
 72      )
 73      color_hexes: list[str] | None = Field(
 74          default=None,
 75          description="For multi-color filaments. List of hex color codes in hex format.",
 76          examples=[["2c3232", "5f5f5f"]],
 77      )
 78      extruder_temp: int | None = Field(default=None, description="Extruder/nozzle temperature in °C.", examples=[210])
 79      bed_temp: int | None = Field(default=None, description="Bed temperature in °C.", examples=[50])
 80      finish: Finish | None = Field(default=None, description="Finish of the filament.", examples=[Finish.MATTE])
 81      multi_color_direction: MultiColorDirection | None = Field(
 82          default=None,
 83          description="Direction of multi-color filaments.",
 84          examples=[MultiColorDirection.COAXIAL],
 85      )
 86      pattern: Pattern | None = Field(default=None, description="Pattern of the filament.", examples=[Pattern.MARBLE])
 87      translucent: bool = Field(default=False, description="Whether the filament is translucent.")
 88      glow: bool = Field(default=False, description="Whether the filament is glow-in-the-dark.")
 89  
 90  
 91  class ExternalFilamentsFile(RootModel):
 92      root: list[ExternalFilament]
 93  
 94      def __iter__(self) -> Iterator[ExternalFilament]:
 95          """Iterate over the filaments."""
 96          return iter(self.root)
 97  
 98      def __getitem__(self, index: int) -> ExternalFilament:
 99          """Get a specific filament by index."""
100          return self.root[index]
101  
102  
103  class ExternalMaterial(BaseModel):
104      material: str = Field(examples=["PLA"])
105      density: float = Field(examples=[1.24])
106      extruder_temp: int | None = Field(default=None, description="Extruder/nozzle temperature in °C.", examples=[210])
107      bed_temp: int | None = Field(default=None, description="Bed temperature in °C.", examples=[50])
108  
109  
110  class ExternalMaterialsFile(RootModel):
111      root: list[ExternalMaterial]
112  
113      def __iter__(self) -> Iterator[ExternalMaterial]:
114          """Iterate over the materials."""
115          return iter(self.root)
116  
117      def __getitem__(self, index: int) -> ExternalMaterial:
118          """Get a specific material by index."""
119          return self.root[index]
120  
121  
122  def get_external_db_url() -> str:
123      """Get the external database URL from environment variables. Defaults to DEFAULT_EXTERNAL_DB_URL."""
124      return os.getenv("EXTERNAL_DB_URL", DEFAULT_EXTERNAL_DB_URL)
125  
126  
127  def get_external_db_sync_interval() -> int:
128      """Get the external database sync interval from environment variables. Defaults to DEFAULT_SYNC_INTERVAL."""
129      return int(os.getenv("EXTERNAL_DB_SYNC_INTERVAL", DEFAULT_SYNC_INTERVAL))
130  
131  
132  async def _download_file(url: str) -> bytes:
133      """Download a file from a URL and return the contents as a string.
134  
135      Uses a file-based cache.
136      """
137      async with hishel.AsyncCacheClient(storage=cache_storage, controller=controller) as client:
138          response = await client.get(url)
139          response.raise_for_status()
140          return response.read()
141  
142  
143  def _parse_filaments_from_bytes(data: bytes) -> ExternalFilamentsFile:
144      """Parse a bytes string into a list of ExternalFilament objects."""
145      return ExternalFilamentsFile.parse_raw(data)
146  
147  
148  def _parse_materials_from_bytes(data: bytes) -> ExternalMaterialsFile:
149      """Parse a bytes string into a list of ExternalMaterial objects."""
150      return ExternalMaterialsFile.parse_raw(data)
151  
152  
153  def _write_to_local_cache(filename: str, data: bytes) -> None:
154      """Write data to the local cache."""
155      filecache.update_file(filename, data)
156  
157  
158  def get_filaments_file() -> Path:
159      """Get the path to the filaments file."""
160      return filecache.get_file("filaments.json")
161  
162  
163  def get_materials_file() -> Path:
164      """Get the path to the materials file."""
165      return filecache.get_file("materials.json")
166  
167  
168  async def _sync() -> None:
169      logger.info("Syncing external DB.")
170  
171      url = get_external_db_url()
172  
173      filaments = _parse_filaments_from_bytes(await _download_file(urljoin(url, "filaments.json")))
174      materials = _parse_materials_from_bytes(await _download_file(urljoin(url, "materials.json")))
175  
176      _write_to_local_cache("filaments.json", filaments.json().encode())
177      _write_to_local_cache("materials.json", materials.json().encode())
178  
179      logger.info(
180          "External DB synced. Filaments: %d, Materials: %d",
181          len(filaments.root),
182          len(materials.root),
183      )
184  
185  
186  def schedule_tasks(scheduler: Scheduler) -> None:
187      """Schedule tasks to be executed by the provided scheduler.
188  
189      Args:
190          scheduler: The scheduler to use for scheduling tasks.
191  
192      """
193      if len(get_external_db_url().strip()) == 0:
194          logger.info("External DB URL is empty. Skipping sync.")
195          return
196  
197      logger.info("Scheduling external DB sync.")
198  
199      # Run once on startup
200      scheduler.once(datetime.timedelta(seconds=0), _sync)  # type: ignore[arg-type]
201  
202      sync_interval = get_external_db_sync_interval()
203      if sync_interval > 0:
204          scheduler.cyclic(datetime.timedelta(seconds=DEFAULT_SYNC_INTERVAL), _sync)  # type: ignore[arg-type]
205      else:
206          logger.info("Sync interval is 0, skipping periodic sync of external db.")