Skip to content

Commit

Permalink
convert to strum
Browse files Browse the repository at this point in the history
  • Loading branch information
@rebasedming
rebasedming committed Aug 22, 2024
1 parent 5f58f8f commit 269172f
Show file tree
Hide file tree
Showing 9 changed files with 178 additions and 226 deletions.
244 changes: 101 additions & 143 deletions src/duckdb/csv.rs
View file

Large diffs are not rendered by default.

19 changes: 6 additions & 13 deletions src/duckdb/delta.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,30 +17,23 @@

use anyhow::{anyhow, Result};
use std::collections::HashMap;
use strum::{AsRefStr, EnumIter};

#[derive(EnumIter, AsRefStr, PartialEq, Debug)]
pub enum DeltaOption {
#[strum(serialize = "files")]
Files,
#[strum(serialize = "preserve_casing")]
PreserveCasing,
}

impl DeltaOption {
pub fn as_str(&self) -> &str {
match self {
Self::Files => "files",
Self::PreserveCasing => "preserve_casing",
}
}

pub fn is_required(&self) -> bool {
match self {
Self::Files => true,
Self::PreserveCasing => false,
}
}

pub fn iter() -> impl Iterator<Item = Self> {
[Self::Files, Self::PreserveCasing].into_iter()
}
}

pub fn create_view(
Expand All @@ -51,7 +44,7 @@ pub fn create_view(
let files = format!(
"'{}'",
table_options
.get(DeltaOption::Files.as_str())
.get(DeltaOption::Files.as_ref())
.ok_or_else(|| anyhow!("files option is required"))?
);

Expand All @@ -70,7 +63,7 @@ mod tests {
let table_name = "test";
let schema_name = "main";
let table_options = HashMap::from([(
DeltaOption::Files.as_str().to_string(),
DeltaOption::Files.as_ref().to_string(),
"/data/delta".to_string(),
)]);

Expand Down
23 changes: 8 additions & 15 deletions src/duckdb/iceberg.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,33 +17,26 @@

use anyhow::{anyhow, Result};
use std::collections::HashMap;
use strum::{AsRefStr, EnumIter};

#[derive(EnumIter, AsRefStr, PartialEq, Debug)]
pub enum IcebergOption {
#[strum(serialize = "allow_moved_paths")]
AllowMovedPaths,
#[strum(serialize = "files")]
Files,
#[strum(serialize = "preserve_casing")]
PreserveCasing,
}

impl IcebergOption {
pub fn as_str(&self) -> &str {
match self {
Self::AllowMovedPaths => "allow_moved_paths",
Self::Files => "files",
Self::PreserveCasing => "preserve_casing",
}
}

pub fn is_required(&self) -> bool {
match self {
Self::AllowMovedPaths => false,
Self::Files => true,
Self::PreserveCasing => false,
}
}

pub fn iter() -> impl Iterator<Item = Self> {
[Self::AllowMovedPaths, Self::Files, Self::PreserveCasing].into_iter()
}
}

pub fn create_view(
Expand All @@ -54,12 +47,12 @@ pub fn create_view(
let files = Some(format!(
"'{}'",
table_options
.get(IcebergOption::Files.as_str())
.get(IcebergOption::Files.as_ref())
.ok_or_else(|| anyhow!("files option is required"))?
));

let allow_moved_paths = table_options
.get(IcebergOption::AllowMovedPaths.as_str())
.get(IcebergOption::AllowMovedPaths.as_ref())
.map(|option| format!("allow_moved_paths = {option}"));

let create_iceberg_str = [files, allow_moved_paths]
Expand All @@ -81,7 +74,7 @@ mod tests {
let table_name = "test";
let schema_name = "main";
let table_options = HashMap::from([(
IcebergOption::Files.as_str().to_string(),
IcebergOption::Files.as_ref().to_string(),
"/data/iceberg".to_string(),
)]);

Expand Down
76 changes: 29 additions & 47 deletions src/duckdb/parquet.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,37 +17,34 @@

use anyhow::{anyhow, Result};
use std::collections::HashMap;
use strum::{AsRefStr, EnumIter};

use super::utils;

#[derive(EnumIter, AsRefStr, PartialEq, Debug)]
pub enum ParquetOption {
#[strum(serialize = "binary_as_string")]
BinaryAsString,
#[strum(serialize = "filename")]
FileName,
#[strum(serialize = "file_row_number")]
FileRowNumber,
#[strum(serialize = "files")]
Files,
#[strum(serialize = "hive_partitioning")]
HivePartitioning,
#[strum(serialize = "hive_types")]
HiveTypes,
#[strum(serialize = "hive_types_autocast")]
HiveTypesAutocast,
#[strum(serialize = "preserve_casing")]
PreserveCasing,
#[strum(serialize = "union_by_name")]
UnionByName,
// TODO: EncryptionConfig
}

impl ParquetOption {
pub fn as_str(&self) -> &str {
match self {
Self::BinaryAsString => "binary_as_string",
Self::FileName => "file_name",
Self::FileRowNumber => "file_row_number",
Self::Files => "files",
Self::HivePartitioning => "hive_partitioning",
Self::HiveTypes => "hive_types",
Self::HiveTypesAutocast => "hive_types_autocast",
Self::PreserveCasing => "preserve_casing",
Self::UnionByName => "union_by_name",
}
}

pub fn is_required(&self) -> bool {
match self {
Self::BinaryAsString => false,
Expand All @@ -61,21 +58,6 @@ impl ParquetOption {
Self::UnionByName => false,
}
}

pub fn iter() -> impl Iterator<Item = Self> {
[
Self::BinaryAsString,
Self::FileName,
Self::FileRowNumber,
Self::Files,
Self::HivePartitioning,
Self::HiveTypes,
Self::HiveTypesAutocast,
Self::PreserveCasing,
Self::UnionByName,
]
.into_iter()
}
}

pub fn create_view(
Expand All @@ -85,36 +67,36 @@ pub fn create_view(
) -> Result<String> {
let files = Some(utils::format_csv(
table_options
.get(ParquetOption::Files.as_str())
.get(ParquetOption::Files.as_ref())
.ok_or_else(|| anyhow!("files option is required"))?,
));

let binary_as_string = table_options
.get(ParquetOption::BinaryAsString.as_str())
.get(ParquetOption::BinaryAsString.as_ref())
.map(|option| format!("binary_as_string = {option}"));

let file_name = table_options
.get(ParquetOption::FileName.as_str())
.get(ParquetOption::FileName.as_ref())
.map(|option| format!("filename = {option}"));

let file_row_number = table_options
.get(ParquetOption::FileRowNumber.as_str())
.get(ParquetOption::FileRowNumber.as_ref())
.map(|option| format!("file_row_number = {option}"));

let hive_partitioning = table_options
.get(ParquetOption::HivePartitioning.as_str())
.get(ParquetOption::HivePartitioning.as_ref())
.map(|option| format!("hive_partitioning = {option}"));

let hive_types = table_options
.get(ParquetOption::HiveTypes.as_str())
.get(ParquetOption::HiveTypes.as_ref())
.map(|option| format!("hive_types = {option}"));

let hive_types_autocast = table_options
.get(ParquetOption::HiveTypesAutocast.as_str())
.get(ParquetOption::HiveTypesAutocast.as_ref())
.map(|option| format!("hive_types_autocast = {option}"));

let union_by_name = table_options
.get(ParquetOption::UnionByName.as_str())
.get(ParquetOption::UnionByName.as_ref())
.map(|option| format!("union_by_name = {option}"));

let create_parquet_str = [
Expand Down Expand Up @@ -146,7 +128,7 @@ mod tests {
let schema_name = "main";
let files = "/data/file.parquet";
let table_options =
HashMap::from([(ParquetOption::Files.as_str().to_string(), files.to_string())]);
HashMap::from([(ParquetOption::Files.as_ref().to_string(), files.to_string())]);
let expected = "CREATE VIEW IF NOT EXISTS main.test AS SELECT * FROM read_parquet('/data/file.parquet')";
let actual = create_view(table_name, schema_name, table_options).unwrap();

Expand All @@ -165,7 +147,7 @@ mod tests {
let schema_name = "main";
let files = "/data/file1.parquet, /data/file2.parquet";
let table_options =
HashMap::from([(ParquetOption::Files.as_str().to_string(), files.to_string())]);
HashMap::from([(ParquetOption::Files.as_ref().to_string(), files.to_string())]);

let expected = "CREATE VIEW IF NOT EXISTS main.test AS SELECT * FROM read_parquet(['/data/file1.parquet', '/data/file2.parquet'])";
let actual = create_view(table_name, schema_name, table_options).unwrap();
Expand All @@ -185,35 +167,35 @@ mod tests {
let schema_name = "main";
let table_options = HashMap::from([
(
ParquetOption::Files.as_str().to_string(),
ParquetOption::Files.as_ref().to_string(),
"/data/file.parquet".to_string(),
),
(
ParquetOption::BinaryAsString.as_str().to_string(),
ParquetOption::BinaryAsString.as_ref().to_string(),
"true".to_string(),
),
(
ParquetOption::FileName.as_str().to_string(),
ParquetOption::FileName.as_ref().to_string(),
"false".to_string(),
),
(
ParquetOption::FileRowNumber.as_str().to_string(),
ParquetOption::FileRowNumber.as_ref().to_string(),
"true".to_string(),
),
(
ParquetOption::HivePartitioning.as_str().to_string(),
ParquetOption::HivePartitioning.as_ref().to_string(),
"true".to_string(),
),
(
ParquetOption::HiveTypes.as_str().to_string(),
ParquetOption::HiveTypes.as_ref().to_string(),
"{'release': DATE, 'orders': BIGINT}".to_string(),
),
(
ParquetOption::HiveTypesAutocast.as_str().to_string(),
ParquetOption::HiveTypesAutocast.as_ref().to_string(),
"true".to_string(),
),
(
ParquetOption::UnionByName.as_str().to_string(),
ParquetOption::UnionByName.as_ref().to_string(),
"true".to_string(),
),
]);
Expand Down
22 changes: 22 additions & 0 deletions src/duckdb/secret.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,29 +22,51 @@ use strum::{AsRefStr, EnumIter};
#[derive(EnumIter, AsRefStr, PartialEq, Debug)]
pub enum UserMappingOptions {
// Universal
#[strum(serialize = "type")]
Type,
#[strum(serialize = "provider")]
Provider,
#[strum(serialize = "scope")]
Scope,
#[strum(serialize = "chain")]
Chain,
// S3/GCS/R2
#[strum(serialize = "key_id")]
KeyId,
#[strum(serialize = "secret")]
Secret,
#[strum(serialize = "region")]
Region,
#[strum(serialize = "session_token")]
SessionToken,
#[strum(serialize = "endpoint")]
Endpoint,
#[strum(serialize = "url_style")]
UrlStyle,
#[strum(serialize = "use_ssl")]
UseSsl,
#[strum(serialize = "url_compatibility_mode")]
UrlCompatibilityMode,
#[strum(serialize = "account_id")]
AccountId,
// Azure
#[strum(serialize = "connection_string")]
ConnectionString,
#[strum(serialize = "account_name")]
AccountName,
#[strum(serialize = "tenant_id")]
TenantId,
#[strum(serialize = "client_id")]
ClientId,
#[strum(serialize = "client_secret")]
ClientSecret,
#[strum(serialize = "client_certificate_path")]
ClientCertificatePath,
#[strum(serialize = "http_proxy")]
HttpProxy,
#[strum(serialize = "proxy_user_name")]
ProxyUserName,
#[strum(serialize = "proxy_password")]
ProxyPassword,
}

Expand Down
5 changes: 3 additions & 2 deletions src/fdw/csv.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ use async_std::task;
use duckdb::arrow::array::RecordBatch;
use pgrx::*;
use std::collections::HashMap;
use strum::IntoEnumIterator;
use supabase_wrappers::prelude::*;

use super::base::*;
Expand Down Expand Up @@ -111,14 +112,14 @@ impl ForeignDataWrapper<BaseFdwError> for CsvFdw {
FOREIGN_SERVER_RELATION_ID => {}
FOREIGN_TABLE_RELATION_ID => {
let valid_options: Vec<String> = CsvOption::iter()
.map(|opt| opt.as_str().to_string())
.map(|opt| opt.as_ref().to_string())
.collect();

validate_options(opt_list.clone(), valid_options)?;

for opt in CsvOption::iter() {
if opt.is_required() {
check_options_contain(&opt_list, opt.as_str())?;
check_options_contain(&opt_list, opt.as_ref())?;
}
}
}
Expand Down
5 changes: 3 additions & 2 deletions src/fdw/delta.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ use async_std::task;
use duckdb::arrow::array::RecordBatch;
use pgrx::*;
use std::collections::HashMap;
use strum::IntoEnumIterator;
use supabase_wrappers::prelude::*;

use super::base::*;
Expand Down Expand Up @@ -111,14 +112,14 @@ impl ForeignDataWrapper<BaseFdwError> for DeltaFdw {
FOREIGN_SERVER_RELATION_ID => {}
FOREIGN_TABLE_RELATION_ID => {
let valid_options: Vec<String> = DeltaOption::iter()
.map(|opt| opt.as_str().to_string())
.map(|opt| opt.as_ref().to_string())
.collect();

validate_options(opt_list.clone(), valid_options)?;

for opt in DeltaOption::iter() {
if opt.is_required() {
check_options_contain(&opt_list, opt.as_str())?;
check_options_contain(&opt_list, opt.as_ref())?;
}
}
}
Expand Down
Loading

0 comments on commit 269172f

Please sign in to comment.