asimov_protocol/handle_resolvers/
csv_handle_resolver.rs

1// This is free and unencumbered software released into the public domain.
2
3use crate::{PeerId, ResolveHandle};
4use alloc::string::String;
5use asimov_id::Handle;
6use core::str::FromStr;
7use csv_async::{AsyncReader, AsyncReaderBuilder, StringRecord};
8use futures_lite::{Stream, StreamExt, pin, stream};
9use iroh::EndpointId;
10use std::io::{Error, Result};
11use tokio::fs::File;
12use tokio::io::AsyncReadExt;
13
14#[cfg(not(feature = "std"))]
15use alloc::collections::BTreeSet as Set;
16
17#[cfg(feature = "std")]
18use std::collections::HashSet as Set;
19
20/// A CSV file resolver from ASIMOV handles to peer IDs.
21///
22/// The format of the CSV file is simply `handle,peer_id`.
23/// The records should be sorted for efficient resolution.
24/// The file is read line-by-line, so it can be very large.
25pub struct CsvHandleResolver(AsyncReader<File>);
26
27impl CsvHandleResolver {
28    /// Opens a CSV file for resolving ASIMOV handles.
29    pub async fn open(path: &str) -> std::io::Result<Self> {
30        let file = File::open(path).await?;
31        let reader = AsyncReaderBuilder::new()
32            .has_headers(false)
33            .create_reader(file);
34        Ok(Self::from(reader))
35    }
36
37    pub fn handles(&mut self) -> impl Stream<Item = Result<Handle>> + Send {
38        async_stream::stream! {
39            let mut handles = Set::new();
40            let records = self.records();
41            pin!(records);
42            while let Some(record) = records.next().await {
43                let (handle, _) = record?;
44                if handles.contains(&handle) {
45                    continue; // skip duplicate handles
46                }
47                handles.insert(handle.clone());
48                yield Ok(handle);
49            }
50        }
51    }
52
53    pub fn records(&mut self) -> impl Stream<Item = Result<(Handle, PeerId)>> + Send {
54        async_stream::stream! {
55            self.0.rewind().await?;
56            let mut record = StringRecord::new();
57            while let Ok(true) = self.0.read_record(&mut record).await {
58                let Some(record_handle) = record.get(0) else {
59                    continue; // skip invalid records
60                };
61                let Some(record_endpoint) = record.get(1) else {
62                    continue; // skip invalid records
63                };
64                let Ok(handle) = record_handle.parse::<Handle>() else {
65                    continue; // skip invalid handles
66                };
67                let Ok(endpoint) = record_endpoint.parse::<PeerId>() else {
68                    continue; // skip invalid peer IDs
69                };
70                yield Ok((handle, endpoint));
71            }
72        }
73    }
74}
75
76impl From<AsyncReader<File>> for CsvHandleResolver {
77    fn from(reader: AsyncReader<File>) -> Self {
78        Self(reader)
79    }
80}
81
82impl ResolveHandle for CsvHandleResolver {
83    type Error = std::io::Error;
84
85    /// Resolves a handle into a set of endpoint IDs.
86    fn resolve_handle(
87        &mut self,
88        handle: impl Into<Handle>,
89    ) -> impl Stream<Item = Result<PeerId>> + Send {
90        let handle = handle.into();
91        async_stream::stream! {
92            let mut endpoints = Set::new();
93            let records = self.records();
94            pin!(records);
95            while let Some(record) = records.next().await {
96                let (record_handle, record_endpoint) = record?;
97                if record_handle != handle {
98                    continue; // skip records that don't match
99                }
100                if endpoints.contains(&record_endpoint) {
101                    continue; // skip duplicate endpoints
102                }
103                endpoints.insert(record_endpoint.clone());
104                yield Ok(record_endpoint);
105            }
106        }
107    }
108}
asimov_protocol/handle_resolvers/csv_handle_resolver.rs

asimov_protocol/handle_resolvers/
csv_handle_resolver.rs
