Skip to content

Commit

Permalink
Feature: Adding Bulk Company Enrich (#24)
Browse files Browse the repository at this point in the history
* Feature: Adding Bulk Company Enrich

This PR updates the SDK to allow for Bulk Company Enrich.
  • Loading branch information
joshfinnie authored Jan 12, 2024
1 parent aa5e482 commit c0c7ff1
Show file tree
Hide file tree
Showing 3 changed files with 113 additions and 5 deletions.
2 changes: 1 addition & 1 deletion Cargo.toml
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[package]
name = "peopledatalabs"
version = "1.0.1"
version = "1.1.0"
edition = "2021"
description = "A Rust client for the People Data Labs API"
documentation = "https://docs.peopledatalabs.com/docs/rust-sdk"
Expand Down
60 changes: 56 additions & 4 deletions src/api/company.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,14 +2,15 @@ use crate::{
models::{
common::SearchParams,
company::{
CleanCompanyParams, CleanCompanyResponse, CompanyResponse, EnrichCompanyParams,
SearchCompanyResponse,
BulkCompanyEnrichResponse, BulkEnrichCompanyParams, CleanCompanyParams,
CleanCompanyResponse, CompanyResponse, EnrichCompanyParams, SearchCompanyResponse,
},
},
PDLClient, PDLError,
};

pub(crate) static ENRICH_PATH: &str = "/company/enrich";
pub(crate) static COMPANY_BULK_ENRICH_PATH: &str = "/company/enrich/bulk";
pub(crate) static SEARCH_PATH: &str = "/company/search";
pub(crate) static CLEAN_PATH: &str = "/company/clean";

Expand All @@ -26,6 +27,19 @@ impl Company {
.get::<CompanyResponse, EnrichCompanyParams>(ENRICH_PATH, params)
}

/// Bulk Enrich for Companies
/// docs: https://docs.peopledatalabs.com/docs/bulk-company-enrichment-api
pub fn bulk_enrich(
&self,
params: BulkEnrichCompanyParams,
) -> Result<Vec<BulkCompanyEnrichResponse>, PDLError> {
self.client
.post::<Vec<BulkCompanyEnrichResponse>, BulkEnrichCompanyParams>(
COMPANY_BULK_ENRICH_PATH,
params,
)
}

/// Search gives you access to every record in our full Company dataset,
/// which you can filter and segment using a search query.
/// docs: https://docs.peopledatalabs.com/docs/company-search-api
Expand All @@ -47,8 +61,10 @@ impl Company {
#[cfg(test)]
mod tests {
use crate::{
client::PDLClient, BaseParams, CleanCompanyParams, CompanyParams, EnrichCompanyParams,
SearchBaseParams, SearchParams,
client::PDLClient,
models::company::{BulkEnrichCompanyParams, BulkEnrichSingleCompanyParams},
BaseParams, CleanCompanyParams, CompanyParams, EnrichCompanyParams, SearchBaseParams,
SearchParams,
};

use super::Company;
Expand Down Expand Up @@ -78,6 +94,42 @@ mod tests {
assert_eq!(resp.name, Some("google".to_string()));
}

#[test]
fn test_bulk_company_enrich() {
let api_key = std::env::var("PDL_API_KEY").unwrap();
let client = PDLClient::new(&api_key).build();

let company = Company { client };

let mut base_params = BaseParams::default();
base_params.pretty = Some(true);

let mut company_params_1 = CompanyParams::default();
company_params_1.profile = Some("https://www.linkedin.com/company/walmart".to_string());
let mut company_params_2 = CompanyParams::default();
company_params_2.website = Some("google.com".to_string());

let bulk_enrich_single_company_params_1 = BulkEnrichSingleCompanyParams {
params: company_params_1,
};

let bulk_enrich_single_company_params_2 = BulkEnrichSingleCompanyParams {
params: company_params_2,
};

let bulk_enrich_params = BulkEnrichCompanyParams {
requests: vec![
bulk_enrich_single_company_params_1,
bulk_enrich_single_company_params_2,
],
};

let resp = company.bulk_enrich(bulk_enrich_params).expect("ERROR");

assert_eq!(resp[0].status, 200);
assert_eq!(resp[1].status, 200);
}

#[test]
fn test_company_clean() {
let api_key = std::env::var("PDL_API_KEY").unwrap();
Expand Down
56 changes: 56 additions & 0 deletions src/models/company.rs
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,29 @@ pub struct CompanyParams {
pub postal_code: Option<String>,
}

impl CompanyParams {
// Validation function
pub fn validate(&self) -> Result<(), PDLError> {
// Check if at least one field is present
if self.pdl_id.is_some()
|| self.name.is_some()
|| self.website.is_some()
|| self.profile.is_some()
|| self.ticker.is_some()
|| self.location.is_some()
|| self.locality.is_some()
|| self.region.is_some()
|| self.country.is_some()
|| self.street_address.is_some()
|| self.postal_code.is_some()
{
return Ok(());
}

Err(PDLError::ValidationError)
}
}

#[derive(Debug, Serialize, Deserialize)]
pub struct EnrichCompanyParams {
#[serde(flatten)]
Expand All @@ -70,6 +93,32 @@ impl EnrichCompanyParams {
}
}

#[derive(Debug, Deserialize, Serialize)]
pub struct BulkEnrichCompanyParams {
pub requests: Vec<BulkEnrichSingleCompanyParams>,
}

impl BulkEnrichCompanyParams {
pub fn validate(&self) -> Result<(), PDLError> {
for request in &self.requests {
request.validate()?
}
Ok(())
}
}

#[derive(Debug, Deserialize, Serialize)]
pub struct BulkEnrichSingleCompanyParams {
pub params: CompanyParams,
}

impl BulkEnrichSingleCompanyParams {
pub fn validate(&self) -> Result<(), PDLError> {
self.params.validate()?;
Ok(())
}
}

#[derive(Debug, Serialize, Deserialize, Default)]
pub struct CleanCompanyParams {
/// The name of the company
Expand Down Expand Up @@ -245,6 +294,13 @@ pub struct CompanyResponse {
pub likelihood: Option<i32>,
}

#[derive(Debug, Deserialize, Serialize)]
pub struct BulkCompanyEnrichResponse {
pub data: Option<Vec<CompanyResponse>>,
pub status: i32,
pub likelihood: Option<i32>,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct CleanCompanyResponse {
/// See https://docs.peopledatalabs.com/docs/output-response-cleaner-apis#company-cleaner-api-response for more information.
Expand Down

0 comments on commit c0c7ff1

Please sign in to comment.