swiftide_docker_executor/
file_loader.rs

1use std::{borrow::Cow, path::PathBuf};
2
3use codegen::{loader_client::LoaderClient, LoadFilesRequest, NodeResponse};
4use swiftide_core::{indexing::Node, Loader};
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 host_port = &self.executor.host_port;
57        let mut client = tokio::task::block_in_place(|| {
58            Handle::current().block_on(async {
59                LoaderClient::connect(format!("http://127.0.0.1:{host_port}")).await
60            })
61        })
62        .expect("Failed to connect to service");
63
64        let (tx, rx) = tokio::sync::mpsc::channel::<anyhow::Result<Node>>(1000);
65
66        tokio::task::spawn(async move {
67            let stream = match client
68                .load_files(LoadFilesRequest {
69                    root_path: self.path.to_string_lossy().to_string(),
70                    file_extensions: self.extensions,
71                })
72                .await
73            {
74                Ok(stream) => stream,
75                Err(error) => {
76                    tracing::error!(error = ?error, "Failed to load files");
77                    return;
78                }
79            };
80
81            let mut stream = stream.into_inner();
82
83            while let Some(result) = stream.message().await.transpose() {
84                if let Err(e) = tx
85                    .send(
86                        result
87                            .map_err(anyhow::Error::from)
88                            .and_then(TryInto::try_into),
89                    )
90                    .await
91                {
92                    tracing::error!(error = ?e, "error sending node");
93                    break;
94                }
95            }
96        });
97
98        rx.into()
99    }
100}
101
102impl TryInto<Node> for NodeResponse {
103    type Error = anyhow::Error;
104
105    fn try_into(self) -> Result<Node, Self::Error> {
106        Node::builder()
107            .path(self.path)
108            .chunk(self.chunk)
109            .original_size(self.original_size as usize)
110            .build()
111    }
112}