-
-
Notifications
You must be signed in to change notification settings - Fork 6
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
59cd808
commit f776261
Showing
15 changed files
with
306 additions
and
23 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,10 @@ | ||
| [ | ||
| { | ||
| "id": 1, | ||
| "name": "John Doe" | ||
| }, | ||
| { | ||
| "id": 2, | ||
| "name": "Jane Smith" | ||
| } | ||
| ] |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,3 @@ | ||
| {"id": 1,"name": "John Doe"} | ||
| {"id": 2,"name": "Jane Smith"} | ||
|
|
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,112 @@ | ||
| use crate::error::Result; | ||
| use crate::polars::Connection; | ||
| use crate::Error::{ConversionError, InvalidUrl}; | ||
| use async_trait::async_trait; | ||
| use polars::io::SerReader; | ||
| use polars::prelude::{IntoLazy, JsonReader}; | ||
| use polars_sql::SQLContext; | ||
| use std::collections::HashMap; | ||
| use std::fs::File; | ||
| use std::num::NonZeroUsize; | ||
| use url::Url; | ||
|
|
||
| #[derive(Debug)] | ||
| pub struct Driver; | ||
|
|
||
| #[async_trait] | ||
| impl crate::Driver for Driver { | ||
| fn identifier(&self) -> &'static str { | ||
| "json" | ||
| } | ||
|
|
||
| async fn connect( | ||
| &self, | ||
| url: String, | ||
| _password: Option<String>, | ||
| ) -> Result<Box<dyn crate::Connection>> { | ||
| let parsed_url = Url::parse(url.as_str())?; | ||
| let query_parameters: HashMap<String, String> = | ||
| parsed_url.query_pairs().into_owned().collect(); | ||
|
|
||
| // Read Options | ||
| let file_name = query_parameters | ||
| .get("file") | ||
| .ok_or(InvalidUrl("Missing file parameter".to_string()))?; | ||
| let file = File::open(file_name)?; | ||
| let ignore_errors = query_parameters | ||
| .get("ignore_errors") | ||
| .map_or(false, |v| v == "true"); | ||
| let infer_schema_length = match query_parameters.get("infer_schema_length") { | ||
| Some(infer_schema_length) => { | ||
| let length = infer_schema_length | ||
| .parse::<usize>() | ||
| .map_err(|error| ConversionError(error.to_string()))?; | ||
| if length == 0 { | ||
| None | ||
| } else { | ||
| NonZeroUsize::new(length) | ||
| } | ||
| } | ||
| None => NonZeroUsize::new(100), | ||
| }; | ||
|
|
||
| let data_frame = JsonReader::new(file) | ||
| .infer_schema_len(infer_schema_length) | ||
| .set_rechunk(true) | ||
| .with_ignore_errors(ignore_errors) | ||
| .finish()?; | ||
|
|
||
| let table_name = crate::polars::driver::get_table_name(file_name)?; | ||
| let mut context = SQLContext::new(); | ||
| context.register(table_name.as_str(), data_frame.lazy()); | ||
|
|
||
| let connection = Connection::new(url, context).await?; | ||
| Ok(Box::new(connection)) | ||
| } | ||
| } | ||
|
|
||
| #[cfg(test)] | ||
| mod test { | ||
| use crate::{DriverManager, Value}; | ||
|
|
||
| const CRATE_DIRECTORY: &str = env!("CARGO_MANIFEST_DIR"); | ||
|
|
||
| fn database_url() -> String { | ||
| format!("json://?file={CRATE_DIRECTORY}/../datasets/users.json") | ||
| } | ||
|
|
||
| #[tokio::test] | ||
| async fn test_driver_connect() -> anyhow::Result<()> { | ||
| let database_url = database_url(); | ||
| let driver_manager = DriverManager::default(); | ||
| let mut connection = driver_manager.connect(&database_url).await?; | ||
| assert_eq!(&database_url, connection.url()); | ||
| connection.close().await?; | ||
| Ok(()) | ||
| } | ||
|
|
||
| #[tokio::test] | ||
| async fn test_connection_interface() -> anyhow::Result<()> { | ||
| let database_url = database_url(); | ||
| let driver_manager = DriverManager::default(); | ||
| let mut connection = driver_manager.connect(&database_url).await?; | ||
|
|
||
| let mut query_result = connection | ||
| .query("SELECT id, name FROM users ORDER BY id") | ||
| .await?; | ||
|
|
||
| assert_eq!(query_result.columns().await, vec!["id", "name"]); | ||
| assert_eq!( | ||
| query_result.next().await, | ||
| Some(vec![Value::I64(1), Value::String("John Doe".to_string())]) | ||
| ); | ||
| assert_eq!( | ||
| query_result.next().await, | ||
| Some(vec![Value::I64(2), Value::String("Jane Smith".to_string())]) | ||
| ); | ||
| assert!(query_result.next().await.is_none()); | ||
|
|
||
| connection.close().await?; | ||
| Ok(()) | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,3 @@ | ||
| pub mod driver; | ||
|
|
||
| pub use driver::Driver; |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,112 @@ | ||
| use crate::error::Result; | ||
| use crate::polars::Connection; | ||
| use crate::Error::{ConversionError, InvalidUrl}; | ||
| use async_trait::async_trait; | ||
| use polars::io::SerReader; | ||
| use polars::prelude::{IntoLazy, JsonLineReader}; | ||
| use polars_sql::SQLContext; | ||
| use std::collections::HashMap; | ||
| use std::fs::File; | ||
| use std::num::NonZeroUsize; | ||
| use url::Url; | ||
|
|
||
| #[derive(Debug)] | ||
| pub struct Driver; | ||
|
|
||
| #[async_trait] | ||
| impl crate::Driver for Driver { | ||
| fn identifier(&self) -> &'static str { | ||
| "jsonl" | ||
| } | ||
|
|
||
| async fn connect( | ||
| &self, | ||
| url: String, | ||
| _password: Option<String>, | ||
| ) -> Result<Box<dyn crate::Connection>> { | ||
| let parsed_url = Url::parse(url.as_str())?; | ||
| let query_parameters: HashMap<String, String> = | ||
| parsed_url.query_pairs().into_owned().collect(); | ||
|
|
||
| // Read Options | ||
| let file_name = query_parameters | ||
| .get("file") | ||
| .ok_or(InvalidUrl("Missing file parameter".to_string()))?; | ||
| let file = File::open(file_name)?; | ||
| let ignore_errors = query_parameters | ||
| .get("ignore_errors") | ||
| .map_or(false, |v| v == "true"); | ||
| let infer_schema_length = match query_parameters.get("infer_schema_length") { | ||
| Some(infer_schema_length) => { | ||
| let length = infer_schema_length | ||
| .parse::<usize>() | ||
| .map_err(|error| ConversionError(error.to_string()))?; | ||
| if length == 0 { | ||
| None | ||
| } else { | ||
| NonZeroUsize::new(length) | ||
| } | ||
| } | ||
| None => NonZeroUsize::new(100), | ||
| }; | ||
|
|
||
| let data_frame = JsonLineReader::new(file) | ||
| .infer_schema_len(infer_schema_length) | ||
| .set_rechunk(true) | ||
| .with_ignore_errors(ignore_errors) | ||
| .finish()?; | ||
|
|
||
| let table_name = crate::polars::driver::get_table_name(file_name)?; | ||
| let mut context = SQLContext::new(); | ||
| context.register(table_name.as_str(), data_frame.lazy()); | ||
|
|
||
| let connection = Connection::new(url, context).await?; | ||
| Ok(Box::new(connection)) | ||
| } | ||
| } | ||
|
|
||
| #[cfg(test)] | ||
| mod test { | ||
| use crate::{DriverManager, Value}; | ||
|
|
||
| const CRATE_DIRECTORY: &str = env!("CARGO_MANIFEST_DIR"); | ||
|
|
||
| fn database_url() -> String { | ||
| format!("jsonl://?file={CRATE_DIRECTORY}/../datasets/users.jsonl") | ||
| } | ||
|
|
||
| #[tokio::test] | ||
| async fn test_driver_connect() -> anyhow::Result<()> { | ||
| let database_url = database_url(); | ||
| let driver_manager = DriverManager::default(); | ||
| let mut connection = driver_manager.connect(&database_url).await?; | ||
| assert_eq!(&database_url, connection.url()); | ||
| connection.close().await?; | ||
| Ok(()) | ||
| } | ||
|
|
||
| #[tokio::test] | ||
| async fn test_connection_interface() -> anyhow::Result<()> { | ||
| let database_url = database_url(); | ||
| let driver_manager = DriverManager::default(); | ||
| let mut connection = driver_manager.connect(&database_url).await?; | ||
|
|
||
| let mut query_result = connection | ||
| .query("SELECT id, name FROM users ORDER BY id") | ||
| .await?; | ||
|
|
||
| assert_eq!(query_result.columns().await, vec!["id", "name"]); | ||
| assert_eq!( | ||
| query_result.next().await, | ||
| Some(vec![Value::I64(1), Value::String("John Doe".to_string())]) | ||
| ); | ||
| assert_eq!( | ||
| query_result.next().await, | ||
| Some(vec![Value::I64(2), Value::String("Jane Smith".to_string())]) | ||
| ); | ||
| assert!(query_result.next().await.is_none()); | ||
|
|
||
| connection.close().await?; | ||
| Ok(()) | ||
| } | ||
| } |
Oops, something went wrong.