Skip to content

Commit

Permalink
Simplified DLP constructors and required options
Browse files Browse the repository at this point in the history
  • Loading branch information
abdolence committed Aug 10, 2024
1 parent 38d496a commit ae8245f
Show file tree
Hide file tree
Showing 8 changed files with 92 additions and 79 deletions.
11 changes: 7 additions & 4 deletions src/args.rs
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
use crate::common_types::GcpProjectId;
use crate::errors::AppError;
use crate::redacters::{
GcpDlpRedacterOptions, GeminiLlmModelName, OpenAiLlmApiKey, OpenAiModelName, RedacterOptions,
RedacterProviderOptions,
GcpDlpRedacterOptions, GeminiLlmModelName, OpenAiLlmApiKey, OpenAiModelName,
RedacterBaseOptions, RedacterOptions, RedacterProviderOptions,
};
use clap::*;
use std::fmt::Display;
Expand Down Expand Up @@ -213,12 +213,15 @@ impl TryInto<RedacterOptions> for RedacterArgs {
message: "Redacter type is required".to_string(),
}),
}?;
Ok(RedacterOptions {
provider_options,
let base_options = RedacterBaseOptions {
allow_unsupported_copies: self.allow_unsupported_copies,
csv_headers_disable: self.csv_headers_disable,
csv_delimiter: self.csv_delimiter.map(|c| c as u8),
sampling_size: self.sampling_size,
};
Ok(RedacterOptions {
provider_options,
base_options,
})
}
}
14 changes: 8 additions & 6 deletions src/commands/copy_command.rs
Original file line number Diff line number Diff line change
Expand Up @@ -51,12 +51,14 @@ pub async fn command_copy(
} else {
bold_style.clone().red().apply_to("✗ No".to_string())
};
let sampling_output =
if let Some(ref sampling_size) = redacter_options.as_ref().and_then(|o| o.sampling_size) {
Style::new().apply_to(format!("{} bytes.", sampling_size))
} else {
Style::new().dim().apply_to("-".to_string())
};
let sampling_output = if let Some(ref sampling_size) = redacter_options
.as_ref()
.and_then(|o| o.base_options.sampling_size)
{
Style::new().apply_to(format!("{} bytes.", sampling_size))
} else {
Style::new().dim().apply_to("-".to_string())
};
term.write_line(
format!(
"Copying from {} to {}.\nRedacting: {}.\nSampling: {}\n",
Expand Down
21 changes: 8 additions & 13 deletions src/redacters/aws_comprehend.rs
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
use crate::errors::AppError;
use crate::filesystems::FileSystemRef;
use crate::redacters::{
RedactSupportedOptions, Redacter, RedacterDataItem, RedacterDataItemContent, RedacterOptions,
Redacters,
RedactSupportedOptions, Redacter, RedacterBaseOptions, RedacterDataItem,
RedacterDataItemContent, Redacters,
};
use crate::reporter::AppReporter;
use crate::AppResult;
Expand All @@ -17,13 +17,13 @@ pub struct AwsComprehendRedacterOptions {
#[derive(Clone)]
pub struct AwsComprehendRedacter<'a> {
client: aws_sdk_comprehend::Client,
redacter_options: RedacterOptions,
base_options: RedacterBaseOptions,
reporter: &'a AppReporter<'a>,
}

impl<'a> AwsComprehendRedacter<'a> {
pub async fn new(
redacter_options: RedacterOptions,
base_options: RedacterBaseOptions,
aws_dlp_options: AwsComprehendRedacterOptions,
reporter: &'a AppReporter<'a>,
) -> AppResult<Self> {
Expand All @@ -35,7 +35,7 @@ impl<'a> AwsComprehendRedacter<'a> {
let client = aws_sdk_comprehend::Client::new(&shared_config);
Ok(Self {
client,
redacter_options,
base_options,
reporter,
})
}
Expand Down Expand Up @@ -114,8 +114,8 @@ impl<'a> Redacter for AwsComprehendRedacter<'a> {
})
}

fn options(&self) -> &RedacterOptions {
&self.redacter_options
fn options(&self) -> &RedacterBaseOptions {
&self.base_options
}
}

Expand All @@ -142,12 +142,7 @@ mod tests {
let content = RedacterDataItemContent::Value(test_content.to_string());
let input = RedacterDataItem { file_ref, content };

let redacter_options = RedacterOptions {
provider_options: RedacterProviderOptions::AwsComprehend(
AwsComprehendRedacterOptions {
region: Some(Region::new(test_aws_region.clone())),
},
),
let redacter_options = RedacterBaseOptions {
allow_unsupported_copies: false,
csv_headers_disable: false,
csv_delimiter: None,
Expand Down
19 changes: 8 additions & 11 deletions src/redacters/gcp_dlp.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,8 @@ use crate::common_types::GcpProjectId;
use crate::errors::AppError;
use crate::filesystems::FileSystemRef;
use crate::redacters::{
RedactSupportedOptions, Redacter, RedacterDataItem, RedacterDataItemContent, RedacterOptions,
Redacters,
RedactSupportedOptions, Redacter, RedacterBaseOptions, RedacterDataItem,
RedacterDataItemContent, Redacters,
};
use crate::reporter::AppReporter;
use crate::AppResult;
Expand All @@ -16,7 +16,7 @@ use rvstruct::ValueStruct;
#[derive(Clone)]
pub struct GcpDlpRedacter<'a> {
client: GoogleApi<DlpServiceClient<GoogleAuthMiddleware>>,
redacter_options: RedacterOptions,
base_options: RedacterBaseOptions,
gcp_dlp_options: GcpDlpRedacterOptions,
reporter: &'a AppReporter<'a>,
}
Expand Down Expand Up @@ -50,7 +50,7 @@ impl<'a> GcpDlpRedacter<'a> {
"ENCRYPTION_KEY",
];
pub async fn new(
redacter_options: RedacterOptions,
base_options: RedacterBaseOptions,
gcp_dlp_options: GcpDlpRedacterOptions,
reporter: &'a AppReporter<'a>,
) -> AppResult<Self> {
Expand All @@ -59,7 +59,7 @@ impl<'a> GcpDlpRedacter<'a> {
.await?;
Ok(GcpDlpRedacter {
client,
redacter_options,
base_options,
gcp_dlp_options,
reporter,
})
Expand Down Expand Up @@ -231,8 +231,8 @@ impl<'a> Redacter for GcpDlpRedacter<'a> {
)
}

fn options(&self) -> &RedacterOptions {
&self.redacter_options
fn options(&self) -> &RedacterBaseOptions {
&self.base_options
}
}

Expand Down Expand Up @@ -402,10 +402,7 @@ mod tests {
let content = RedacterDataItemContent::Value(test_content.to_string());
let input = RedacterDataItem { file_ref, content };

let redacter_options = RedacterOptions {
provider_options: RedacterProviderOptions::GcpDlp(GcpDlpRedacterOptions {
project_id: GcpProjectId::new(test_gcp_project_id.clone()),
}),
let redacter_options = RedacterBaseOptions {
allow_unsupported_copies: false,
csv_headers_disable: false,
csv_delimiter: None,
Expand Down
20 changes: 8 additions & 12 deletions src/redacters/gemini_llm.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,8 @@ use crate::common_types::GcpProjectId;
use crate::errors::AppError;
use crate::filesystems::FileSystemRef;
use crate::redacters::{
RedactSupportedOptions, Redacter, RedacterDataItem, RedacterDataItemContent, RedacterOptions,
Redacters,
RedactSupportedOptions, Redacter, RedacterBaseOptions, RedacterDataItem,
RedacterDataItemContent, Redacters,
};
use crate::reporter::AppReporter;
use crate::AppResult;
Expand All @@ -24,7 +24,7 @@ pub struct GeminiLlmModelName(String);
#[derive(Clone)]
pub struct GeminiLlmRedacter<'a> {
client: GoogleApi<GenerativeServiceClient<GoogleAuthMiddleware>>,
redacter_options: RedacterOptions,
base_options: RedacterBaseOptions,
gemini_llm_options: crate::redacters::GeminiLlmRedacterOptions,
reporter: &'a AppReporter<'a>,
}
Expand All @@ -33,7 +33,7 @@ impl<'a> GeminiLlmRedacter<'a> {
const DEFAULT_GEMINI_MODEL: &'static str = "models/gemini-1.5-flash";

pub async fn new(
redacter_options: RedacterOptions,
base_options: RedacterBaseOptions,
gemini_llm_options: GeminiLlmRedacterOptions,
reporter: &'a AppReporter<'a>,
) -> AppResult<Self> {
Expand All @@ -47,7 +47,7 @@ impl<'a> GeminiLlmRedacter<'a> {
).await?;
Ok(GeminiLlmRedacter {
client,
redacter_options,
base_options,
gemini_llm_options,
reporter,
})
Expand Down Expand Up @@ -199,8 +199,8 @@ impl<'a> Redacter for GeminiLlmRedacter<'a> {
})
}

fn options(&self) -> &RedacterOptions {
&self.redacter_options
fn options(&self) -> &RedacterBaseOptions {
&self.base_options
}
}

Expand Down Expand Up @@ -228,11 +228,7 @@ mod tests {
let content = RedacterDataItemContent::Value(test_content.to_string());
let input = RedacterDataItem { file_ref, content };

let redacter_options = RedacterOptions {
provider_options: RedacterProviderOptions::GeminiLlm(GeminiLlmRedacterOptions {
project_id: GcpProjectId::new(test_gcp_project_id.clone()),
gemini_model: None,
}),
let redacter_options = RedacterBaseOptions {
allow_unsupported_copies: false,
csv_headers_disable: false,
csv_delimiter: None,
Expand Down
46 changes: 37 additions & 9 deletions src/redacters/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,11 @@ pub enum Redacters<'a> {
#[derive(Debug, Clone)]
pub struct RedacterOptions {
pub provider_options: RedacterProviderOptions,
pub base_options: RedacterBaseOptions,
}

#[derive(Debug, Clone)]
pub struct RedacterBaseOptions {
pub allow_unsupported_copies: bool,
pub csv_headers_disable: bool,
pub csv_delimiter: Option<u8>,
Expand Down Expand Up @@ -87,21 +92,44 @@ impl<'a> Redacters<'a> {
) -> AppResult<Self> {
match redacter_options.provider_options {
RedacterProviderOptions::GcpDlp(ref options) => Ok(Redacters::GcpDlp(
GcpDlpRedacter::new(redacter_options.clone(), options.clone(), reporter).await?,
GcpDlpRedacter::new(
redacter_options.base_options.clone(),
options.clone(),
reporter,
)
.await?,
)),
RedacterProviderOptions::AwsComprehend(ref options) => Ok(Redacters::AwsComprehendDlp(
AwsComprehendRedacter::new(redacter_options.clone(), options.clone(), reporter)
.await?,
AwsComprehendRedacter::new(
redacter_options.base_options.clone(),
options.clone(),
reporter,
)
.await?,
)),
RedacterProviderOptions::MsPresidio(ref options) => Ok(Redacters::MsPresidio(
MsPresidioRedacter::new(redacter_options.clone(), options.clone(), reporter)
.await?,
MsPresidioRedacter::new(
redacter_options.base_options.clone(),
options.clone(),
reporter,
)
.await?,
)),
RedacterProviderOptions::GeminiLlm(ref options) => Ok(Redacters::GeminiLlm(
GeminiLlmRedacter::new(redacter_options.clone(), options.clone(), reporter).await?,
GeminiLlmRedacter::new(
redacter_options.base_options.clone(),
options.clone(),
reporter,
)
.await?,
)),
RedacterProviderOptions::OpenAiLlm(ref options) => Ok(Redacters::OpenAiLlm(
OpenAiLlmRedacter::new(redacter_options.clone(), options.clone(), reporter).await?,
OpenAiLlmRedacter::new(
redacter_options.base_options.clone(),
options.clone(),
reporter,
)
.await?,
)),
}
}
Expand Down Expand Up @@ -146,7 +174,7 @@ pub trait Redacter {
file_ref: &FileSystemRef,
) -> AppResult<RedactSupportedOptions>;

fn options(&self) -> &RedacterOptions;
fn options(&self) -> &RedacterBaseOptions;
}

impl<'a> Redacter for Redacters<'a> {
Expand Down Expand Up @@ -175,7 +203,7 @@ impl<'a> Redacter for Redacters<'a> {
}
}

fn options(&self) -> &RedacterOptions {
fn options(&self) -> &RedacterBaseOptions {
match self {
Redacters::GcpDlp(redacter) => redacter.options(),
Redacters::AwsComprehendDlp(redacter) => redacter.options(),
Expand Down
20 changes: 8 additions & 12 deletions src/redacters/ms_presidio.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,8 @@ use url::Url;
use crate::errors::AppError;
use crate::filesystems::FileSystemRef;
use crate::redacters::{
RedactSupportedOptions, Redacter, RedacterDataItem, RedacterDataItemContent, RedacterOptions,
Redacters,
RedactSupportedOptions, Redacter, RedacterBaseOptions, RedacterDataItem,
RedacterDataItemContent, Redacters,
};
use crate::reporter::AppReporter;
use crate::AppResult;
Expand All @@ -21,7 +21,7 @@ pub struct MsPresidioRedacterOptions {
pub struct MsPresidioRedacter<'a> {
client: reqwest::Client,
ms_presidio_options: MsPresidioRedacterOptions,
redacter_options: RedacterOptions,
base_options: RedacterBaseOptions,
reporter: &'a AppReporter<'a>,
}

Expand All @@ -44,15 +44,15 @@ impl<'a> MsPresidioRedacter<'a> {
const DISALLOW_ENTITY_TYPES: [&'static str; 1] = ["US_DRIVER_LICENSE"];

pub async fn new(
redacter_options: RedacterOptions,
base_options: RedacterBaseOptions,
ms_presidio_options: MsPresidioRedacterOptions,
reporter: &'a AppReporter<'a>,
) -> AppResult<Self> {
let client = reqwest::Client::new();
Ok(Self {
client,
ms_presidio_options,
redacter_options,
base_options,
reporter,
})
}
Expand Down Expand Up @@ -215,8 +215,8 @@ impl<'a> Redacter for MsPresidioRedacter<'a> {
})
}

fn options(&self) -> &RedacterOptions {
&self.redacter_options
fn options(&self) -> &RedacterBaseOptions {
&self.base_options
}
}

Expand Down Expand Up @@ -249,11 +249,7 @@ mod tests {
let content = RedacterDataItemContent::Value(test_content.to_string());
let input = RedacterDataItem { file_ref, content };

let redacter_options = RedacterOptions {
provider_options: RedacterProviderOptions::MsPresidio(MsPresidioRedacterOptions {
text_analyze_url: Some(test_analyze_url.clone()),
image_redact_url: None,
}),
let redacter_options = RedacterBaseOptions {
allow_unsupported_copies: false,
csv_headers_disable: false,
csv_delimiter: None,
Expand Down
Loading

0 comments on commit ae8245f

Please sign in to comment.