swiftide_docker_executor/
file_loader.rs

1use std::{borrow::Cow, path::PathBuf};
2
3use codegen::{LoadFilesRequest, NodeResponse, loader_client::LoaderClient};
4use swiftide_core::{Loader, indexing::Node};
5use tokio::runtime::Handle;
6
7use crate::RunningDockerExecutor;
8
9pub mod codegen {
10    tonic::include_proto!("loader");
11}
12
13#[derive(Debug, Clone)]
14pub struct FileLoader<'a> {
15    path: PathBuf,
16    extensions: Vec<String>,
17    executor: Cow<'a, RunningDockerExecutor>,
18}
19
20impl RunningDockerExecutor {
21    /// Creates an owned file loader from the executor. If needed it is safe to clone the executor.
22    ///
23    /// The loader can be used with a swiftide indexing pipeline.
24    pub fn into_file_loader<V: IntoIterator<Item = T>, T: Into<String>>(
25        self,
26        path: impl Into<PathBuf>,
27        extensions: V,
28    ) -> FileLoader<'static> {
29        let path = path.into();
30        let extensions = extensions.into_iter().map(Into::into).collect::<Vec<_>>();
31        FileLoader {
32            path,
33            extensions,
34            executor: Cow::Owned(self),
35        }
36    }
37
38    /// Creates a borrowed file loader from the executor.
39    pub fn as_file_loader<'a, V: IntoIterator<Item = T>, T: Into<String>>(
40        &'a self,
41        path: impl Into<PathBuf>,
42        extensions: V,
43    ) -> FileLoader<'a> {
44        let path = path.into();
45        let extensions = extensions.into_iter().map(Into::into).collect::<Vec<_>>();
46        FileLoader {
47            path,
48            extensions,
49            executor: Cow::Borrowed(self),
50        }
51    }
52}
53
54impl Loader for FileLoader<'_> {
55    fn into_stream(self) -> swiftide_core::indexing::IndexingStream {
56        let container_ip = &self.executor.container_ip;
57        let container_port = &self.executor.container_port;
58        let mut client = tokio::task::block_in_place(|| {
59            Handle::current().block_on(async {
60                LoaderClient::connect(format!("http://{container_ip}:{container_port}")).await
61            })
62        })
63        .expect("Failed to connect to service");
64
65        let (tx, rx) = tokio::sync::mpsc::channel::<anyhow::Result<Node>>(1000);
66
67        tokio::task::spawn(async move {
68            let stream = match client
69                .load_files(LoadFilesRequest {
70                    root_path: self.path.to_string_lossy().to_string(),
71                    file_extensions: self.extensions,
72                })
73                .await
74            {
75                Ok(stream) => stream,
76                Err(error) => {
77                    tracing::error!(error = ?error, "Failed to load files");
78                    return;
79                }
80            };
81
82            let mut stream = stream.into_inner();
83
84            while let Some(result) = stream.message().await.transpose() {
85                if let Err(e) = tx
86                    .send(
87                        result
88                            .map_err(anyhow::Error::from)
89                            .and_then(TryInto::try_into),
90                    )
91                    .await
92                {
93                    tracing::error!(error = ?e, "error sending node");
94                    break;
95                }
96            }
97        });
98
99        rx.into()
100    }
101}
102
103impl TryInto<Node> for NodeResponse {
104    type Error = anyhow::Error;
105
106    fn try_into(self) -> Result<Node, Self::Error> {
107        Node::builder()
108            .path(self.path)
109            .chunk(self.chunk)
110            .original_size(self.original_size as usize)
111            .build()
112    }
113}