Skip to content

Commit

Permalink
feat: reject invalid timestamp ranges in copy statement (#3623)
Browse files Browse the repository at this point in the history
* chore: reject invalid timestamp ranges in copy statement

* tests: add unit tests
  • Loading branch information
v0y4g3r authored Apr 1, 2024
1 parent d6b2d1d commit 5e24448
Show file tree
Hide file tree
Showing 3 changed files with 43 additions and 17 deletions.
10 changes: 9 additions & 1 deletion src/operator/src/error.rs
Original file line number Diff line number Diff line change
Expand Up @@ -534,6 +534,13 @@ pub enum Error {
source: session::session_config::Error,
location: Location,
},

#[snafu(display("Invalid timestamp range, start: `{}`, end: `{}`", start, end))]
InvalidTimestampRange {
start: String,
end: String,
location: Location,
},
}

pub type Result<T> = std::result::Result<T, Error>;
Expand Down Expand Up @@ -658,7 +665,8 @@ impl ErrorExt for Error {
| Error::DdlWithMultiSchemas { .. }
| Error::EmptyDdlExpr { .. }
| Error::InvalidPartitionRule { .. }
| Error::ParseSqlValue { .. } => StatusCode::InvalidArguments,
| Error::ParseSqlValue { .. }
| Error::InvalidTimestampRange { .. } => StatusCode::InvalidArguments,

Error::CreateLogicalTables { .. } => StatusCode::Unexpected,
}
Expand Down
2 changes: 2 additions & 0 deletions src/operator/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,8 @@
// See the License for the specific language governing permissions and
// limitations under the License.

#![feature(assert_matches)]

pub mod delete;
pub mod error;
pub mod expr_factory;
Expand Down
48 changes: 32 additions & 16 deletions src/operator/src/statement.rs
Original file line number Diff line number Diff line change
Expand Up @@ -348,7 +348,12 @@ fn timestamp_range_from_option_map(
let start_timestamp = extract_timestamp(options, COPY_DATABASE_TIME_START_KEY, query_ctx)?;
let end_timestamp = extract_timestamp(options, COPY_DATABASE_TIME_END_KEY, query_ctx)?;
let time_range = match (start_timestamp, end_timestamp) {
(Some(start), Some(end)) => TimestampRange::new(start, end),
(Some(start), Some(end)) => Some(TimestampRange::new(start, end).with_context(|| {
error::InvalidTimestampRangeSnafu {
start: start.to_iso8601_string(),
end: end.to_iso8601_string(),
}
})?),
(Some(start), None) => Some(TimestampRange::from_start(start)),
(None, Some(end)) => Some(TimestampRange::until_end(end, false)), // exclusive end
(None, None) => None,
Expand Down Expand Up @@ -391,41 +396,52 @@ fn idents_to_full_database_name(

#[cfg(test)]
mod tests {
use std::assert_matches::assert_matches;
use std::collections::HashMap;
use std::sync::Arc;

use common_time::range::TimestampRange;
use common_time::{Timestamp, Timezone};
use session::context::QueryContextBuilder;
use sql::statements::OptionMap;

use crate::error;
use crate::statement::copy_database::{
COPY_DATABASE_TIME_END_KEY, COPY_DATABASE_TIME_START_KEY,
};
use crate::statement::timestamp_range_from_option_map;

#[test]
fn test_timestamp_range_from_option_map() {
fn check_timestamp_range((start, end): (&str, &str)) -> error::Result<Option<TimestampRange>> {
let query_ctx = QueryContextBuilder::default()
.timezone(Arc::new(Timezone::from_tz_string("Asia/Shanghai").unwrap()))
.build();
let map = OptionMap::from(
[
(
COPY_DATABASE_TIME_START_KEY.to_string(),
"2022-04-11 08:00:00".to_string(),
),
(
COPY_DATABASE_TIME_END_KEY.to_string(),
"2022-04-11 16:00:00".to_string(),
),
(COPY_DATABASE_TIME_START_KEY.to_string(), start.to_string()),
(COPY_DATABASE_TIME_END_KEY.to_string(), end.to_string()),
]
.into_iter()
.collect::<HashMap<_, _>>(),
);
let range = timestamp_range_from_option_map(&map, &query_ctx)
.unwrap()
.unwrap();
assert_eq!(Timestamp::new_second(1649635200), range.start().unwrap());
assert_eq!(Timestamp::new_second(1649664000), range.end().unwrap());
timestamp_range_from_option_map(&map, &query_ctx)
}

#[test]
fn test_timestamp_range_from_option_map() {
assert_eq!(
Some(
TimestampRange::new(
Timestamp::new_second(1649635200),
Timestamp::new_second(1649664000),
)
.unwrap(),
),
check_timestamp_range(("2022-04-11 08:00:00", "2022-04-11 16:00:00"),).unwrap()
);

assert_matches!(
check_timestamp_range(("2022-04-11 08:00:00", "2022-04-11 07:00:00")).unwrap_err(),
error::Error::InvalidTimestampRange { .. }
);
}
}

0 comments on commit 5e24448

Please sign in to comment.