pybgpkitstream 0.1.2__tar.gz → 0.1.4__tar.gz

{pybgpkitstream-0.1.2 → pybgpkitstream-0.1.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: pybgpkitstream
-Version: 0.1.2
+Version: 0.1.4
 Summary: Drop-in replacement for PyBGPStream using BGPKIT
 Author: JustinLoye
 Author-email: JustinLoye <jloye@iij.ad.jp>

{pybgpkitstream-0.1.2 → pybgpkitstream-0.1.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "pybgpkitstream"
-version = "0.1.2"
+version = "0.1.4"
 description = "Drop-in replacement for PyBGPStream using BGPKIT"
 readme = "README.md"
 authors = [

{pybgpkitstream-0.1.2 → pybgpkitstream-0.1.4}/src/pybgpkitstream/bgpkitstream.py

@@ -9,6 +9,7 @@ from heapq import merge
 from operator import itemgetter
 import binascii
 import logging
+from tempfile import TemporaryDirectory
 
 import aiohttp
 import bgpkit
@@ -44,24 +45,39 @@ def crc32(input_str: str):
     return f"{crc:08x}"
 
 
+class Directory:
+    """Permanent directory that mimics TemporaryDirectory interface."""
+
+    def __init__(self, path):
+        self.name = str(path)
+
+    def cleanup(self):
+        """No-op cleanup for permanent directories."""
+        pass
+
+
 class BGPKITStream:
     def __init__(
         self,
-        ts_start: datetime.datetime,
-        ts_end: datetime.datetime,
+        ts_start: float,
+        ts_end: float,
         collector_id: str,
         data_type: list[Literal["update", "rib"]],
         cache_dir: str | None,
         filters: dict = {},
         max_concurrent_downloads: int = 10,
+        chunk_time: float | None = datetime.timedelta(hours=2).seconds,
     ):
         self.ts_start = ts_start
         self.ts_end = ts_end
         self.collector_id = collector_id
         self.data_type = data_type
-        self.cache_dir = cache_dir
+        self.cache_dir: Directory | TemporaryDirectory = (
+            Directory(cache_dir) if cache_dir else TemporaryDirectory()
+        )
         self.filters = filters
         self.max_concurrent_downloads = max_concurrent_downloads
+        self.chunk_time = chunk_time
 
         self.broker = bgpkit.Broker()
 
@@ -70,7 +86,6 @@ class BGPKITStream:
         """Generate a cache filename compatible with BGPKIT parser."""
 
         hash_suffix = crc32(url)
-        print(url)
 
         if "updates." in url:
             data_type = "updates"
@@ -140,7 +155,7 @@ class BGPKITStream:
                 for rc, rc_urls in self.urls[data_type].items():
                     for url in rc_urls:
                         filename = self._generate_cache_filename(url)
-                        filepath = os.path.join(self.cache_dir, filename)
+                        filepath = os.path.join(self.cache_dir.name, filename)
 
                         if os.path.exists(filepath):
                             logging.debug(f"{filepath} is a cache hit")
@@ -171,47 +186,79 @@ class BGPKITStream:
         return ((elem.timestamp, elem, is_rib, collector) for elem in iterator)
 
     def __iter__(self) -> Iterator[BGPElement]:
-        self._set_urls()
+        # try/finally to cleanup the fetching cache
+        try:
+            # Manager mode: spawn smaller worker streams to balance fetch/parse
+            if self.chunk_time:
+                current = self.ts_start
+
+                while current < self.ts_end:
+                    chunk_end = min(current + self.chunk_time, self.ts_end)
+
+                    logging.info(
+                        f"Processing chunk: {datetime.datetime.fromtimestamp(current)} "
+                        f"to {datetime.datetime.fromtimestamp(chunk_end)}"
+                    )
+
+                    worker = type(self)(
+                        ts_start=current,
+                        ts_end=chunk_end
+                        - 1,  # remove one second because BGPKIT include border
+                        collector_id=self.collector_id,
+                        data_type=self.data_type,
+                        cache_dir=None,
+                        filters=self.filters,
+                        max_concurrent_downloads=self.max_concurrent_downloads,
+                        chunk_time=None,  # Worker doesn't chunk itself
+                    )
+
+                    yield from worker
+                    current = chunk_end + 1e-7
+
+                return
+
+            self._set_urls()
 
-        if self.cache_dir:
             asyncio.run(self._prefetch_data())
 
-        # One iterator for each data_type * collector combinations
-        # To be merged according to the elements timestamp
-        iterators_to_merge = []
+            # One iterator for each data_type * collector combinations
+            # To be merged according to the elements timestamp
+            iterators_to_merge = []
 
-        for data_type in self.data_type:
-            is_rib = data_type == "rib"
-
-            # Get rib or update files per collector
-            if self.cache_dir:
-                rc_to_urls = self.paths[data_type]
-            else:
-                rc_to_urls = self.urls[data_type]
-
-            # Chain rib or update iterators to get one stream per collector / data_type
-            for rc, urls in rc_to_urls.items():
-                parsers = [bgpkit.Parser(url=url, filters=self.filters) for url in urls]
-
-                chained_iterator = chain.from_iterable(parsers)
-
-                # Add metadata lost by bgpkit for compatibility with pubgpstream
-                iterators_to_merge.append((chained_iterator, is_rib, rc))
-
-        # Make a generator to tag each bgpkit element with metadata
-        # Benefit 1: full compat with pybgpstream
-        # Benefit 2: we give a key easy to access for heapq to merge
-        tagged_iterators = [
-            self._create_tagged_iterator(it, is_rib, rc)
-            for it, is_rib, rc in iterators_to_merge
-        ]
-
-        # Merge and convert to pybgpstream format
-        for timestamp, bgpkit_elem, is_rib, rc in merge(
-            *tagged_iterators, key=itemgetter(0)
-        ):
-            if self.ts_start <= timestamp <= self.ts_end:
-                yield convert_bgpkit_elem(bgpkit_elem, is_rib, rc)
+            for data_type in self.data_type:
+                is_rib = data_type == "rib"
+
+                # Get rib or update files per collector
+                rc_to_paths = self.paths[data_type]
+
+                # Chain rib or update iterators to get one stream per collector / data_type
+                for rc, paths in rc_to_paths.items():
+                    parsers = [
+                        bgpkit.Parser(url=path, filters=self.filters) for path in paths
+                    ]
+
+                    chained_iterator = chain.from_iterable(parsers)
+
+                    # Add metadata lost by bgpkit for compatibility with pubgpstream
+                    iterators_to_merge.append((chained_iterator, is_rib, rc))
+
+            # Make a generator to tag each bgpkit element with metadata
+            # Benefit 1: full compat with pybgpstream
+            # Benefit 2: we give a key easy to access for heapq to merge
+            tagged_iterators = [
+                self._create_tagged_iterator(it, is_rib, rc)
+                for it, is_rib, rc in iterators_to_merge
+            ]
+
+            # Merge and convert to pybgpstream format
+            for timestamp, bgpkit_elem, is_rib, rc in merge(
+                *tagged_iterators, key=itemgetter(0)
+            ):
+                if self.ts_start <= timestamp <= self.ts_end:
+                    yield convert_bgpkit_elem(bgpkit_elem, is_rib, rc)
+
+        finally:
+            self.cache_dir.cleanup()
 
     @classmethod
     def from_config(cls, config: BGPStreamConfig):
@@ -229,4 +276,5 @@ class BGPKITStream:
             max_concurrent_downloads=config.max_concurrent_downloads
             if config.max_concurrent_downloads
             else 10,
+            chunk_time=config.chunk_time.seconds if config.chunk_time else None,
         )

{pybgpkitstream-0.1.2 → pybgpkitstream-0.1.4}/src/pybgpkitstream/bgpstreamconfig.py

@@ -63,4 +63,10 @@ class BGPStreamConfig(BaseModel):
         description="Specifies the directory for caching downloaded files.",
     )
     filters: FilterOptions | None = Field(default=None, description="Optional filters")
-    max_concurrent_downloads: int | None = Field(default=None, description="Maximum concurrent downloads when caching")
+    max_concurrent_downloads: int | None = Field(
+        default=None, description="Maximum concurrent downloads when caching"
+    )
+    chunk_time: datetime.timedelta | None = Field(
+        default=datetime.timedelta(hours=2),
+        description="Interval for the fetch/parse cycle (avoid long prefetch time)",
+    )