File: add option create_only

SeaQL · Sep 21, 2023 · 7f66afc · 7f66afc
1 parent 44585f5
commit 7f66afc
Show file tree

Hide file tree

Showing 3 changed files with 48 additions and 9 deletions.
diff --git a/sea-streamer-file/src/file.rs b/sea-streamer-file/src/file.rs
@@ -148,6 +148,15 @@ impl AsyncFile {
         Self::new_with(id, file).await
     }
 
+    /// Always create a new file. If the file already exists, abort.
+    pub async fn new_w(id: FileId) -> Result<Self, FileErr> {
+        log::debug!("AsyncFile Create ({})", id.path());
+        let mut options = OpenOptions::new();
+        options.write(true).create_new(true);
+        let file = options.open(id.path()).await.map_err(FileErr::IoError)?;
+        Self::new_with(id, file).await
+    }
+
     async fn new_with(id: FileId, file: File) -> Result<Self, FileErr> {
         let size = file_size_of(&file).await?;
         let pos = 0;

diff --git a/sea-streamer-file/src/streamer.rs b/sea-streamer-file/src/streamer.rs
@@ -20,13 +20,22 @@ pub struct FileStreamer {
 
 #[derive(Debug, Clone)]
 pub struct FileConnectOptions {
-    create_if_not_exists: bool,
+    create_file: CreateFileOption,
     end_with_eos: bool,
     beacon_interval: u32,
     file_size_limit: u64,
     prefetch_message: usize,
 }
 
+#[derive(Debug, Clone)]
+enum CreateFileOption {
+    /// File must already exists
+    Never,
+    CreateIfNotExists,
+    /// Fail if the file already exists
+    Always,
+}
+
 #[derive(Debug, Clone)]
 pub struct FileConsumerOptions {
     mode: ConsumerMode,
@@ -88,11 +97,11 @@ impl StreamerTrait for FileStreamer {
             .trim_start_matches("file://")
             .trim_end_matches('/');
         let file_id = FileId::new(path);
-        if options.create_if_not_exists {
-            AsyncFile::new_rw(file_id.clone()).await?;
-        } else {
-            AsyncFile::new_r(file_id.clone()).await?;
-        }
+        match options.create_file {
+            CreateFileOption::Never => AsyncFile::new_r(file_id.clone()).await,
+            CreateFileOption::CreateIfNotExists => AsyncFile::new_rw(file_id.clone()).await,
+            CreateFileOption::Always => AsyncFile::new_w(file_id.clone()).await,
+        }?;
         Ok(Self { file_id, options })
     }
 
@@ -187,7 +196,7 @@ impl ConnectOptionsTrait for FileConnectOptions {
 impl Default for FileConnectOptions {
     fn default() -> Self {
         Self {
-            create_if_not_exists: false,
+            create_file: CreateFileOption::Never,
             end_with_eos: false,
             beacon_interval: DEFAULT_BEACON_INTERVAL,
             file_size_limit: DEFAULT_FILE_SIZE_LIMIT,
@@ -198,11 +207,28 @@ impl Default for FileConnectOptions {
 
 impl FileConnectOptions {
     pub fn create_if_not_exists(&self) -> bool {
-        self.create_if_not_exists
+        matches!(self.create_file, CreateFileOption::CreateIfNotExists)
     }
     /// Default is `false`.
     pub fn set_create_if_not_exists(&mut self, v: bool) -> &mut Self {
-        self.create_if_not_exists = v;
+        if v {
+            self.create_file = CreateFileOption::CreateIfNotExists;
+        } else {
+            self.create_file = CreateFileOption::Never;
+        }
+        self
+    }
+
+    pub fn create_only(&self) -> bool {
+        matches!(self.create_file, CreateFileOption::Always)
+    }
+    /// Always create the file. Fail if already exists. Default is `false`.
+    pub fn set_create_only(&mut self, v: bool) -> &mut Self {
+        if v {
+            self.create_file = CreateFileOption::Always;
+        } else {
+            self.create_file = CreateFileOption::Never;
+        }
         self
     }
 

diff --git a/sea-streamer-runtime/src/file/no_rt_file.rs b/sea-streamer-runtime/src/file/no_rt_file.rs
@@ -110,6 +110,10 @@ impl OpenOptions {
         self
     }
 
+    pub fn create_new(&mut self, _: bool) -> &mut OpenOptions {
+        self
+    }
+
     pub async fn open(&self, _: impl AsRef<Path>) -> Result<File, IoError> {
         Err(IoError::new(ErrorKind::Other, "Please enable a runtime"))
     }