affine_importer 0.1.2

AFFiNE import planning and batching core.
Documentation
use std::{collections::HashMap, path::PathBuf, sync::Arc};

use crate::{ImportBatch, ImportBatchLimits, ImportError, ImportOptions, ImportResult};

#[derive(Debug, Clone)]
pub enum ImportSource {
  FilePath(PathBuf),
  DirectoryPath(PathBuf),
}

pub trait ImportCursor: Send {
  fn total(&self) -> usize;
  fn next_batch(&mut self) -> ImportResult<Option<ImportBatch>>;
}

pub trait ImportProvider: Send + Sync {
  fn format(&self) -> &'static str;

  fn create_cursor(
    &self,
    source: ImportSource,
    options: ImportOptions,
    limits: ImportBatchLimits,
  ) -> ImportResult<Box<dyn ImportCursor>>;
}

#[derive(Default)]
pub struct ImportRegistry {
  providers: HashMap<&'static str, Arc<dyn ImportProvider>>,
}

impl ImportRegistry {
  pub fn new() -> Self {
    Self::default()
  }

  pub fn with_builtin() -> Self {
    let mut registry = Self::new();
    crate::planner::register_builtin_providers(&mut registry);
    registry
  }

  pub fn register<P>(&mut self, provider: P)
  where
    P: ImportProvider + 'static,
  {
    self.providers.insert(provider.format(), Arc::new(provider));
  }

  pub fn create_cursor(
    &self,
    format: &str,
    source: ImportSource,
    options: ImportOptions,
    limits: ImportBatchLimits,
  ) -> ImportResult<Box<dyn ImportCursor>> {
    let provider = self
      .providers
      .get(format)
      .ok_or_else(|| ImportError::UnsupportedFormat(format.to_string()))?;
    provider.create_cursor(source, options, limits)
  }
}