UlfarErl commented on code in PR #7333:
URL: https://github.com/apache/arrow-datafusion/pull/7333#discussion_r1298742405
##########
datafusion/core/src/execution/context.rs:
##########
@@ -2086,6 +2190,92 @@ impl SerializerRegistry for EmptySerializerRegistry {
}
}
+/// Describes which SQL statements can be run.
+///
+/// See [`SessionContext::sql_with_options`] for more details.
+#[derive(Clone, Debug, Copy)]
+pub struct SQLOptions {
+ /// See [`Self::with_allow_ddl`]
+ allow_ddl: bool,
+ /// See [`Self::with_allow_dml`]
+ allow_dml: bool,
+ /// See [`Self::with_allow_statements`]
+ allow_statements: bool,
+}
+
+impl Default for SQLOptions {
+ fn default() -> Self {
+ Self {
+ allow_ddl: true,
+ allow_dml: true,
+ allow_statements: true,
+ }
+ }
+}
+
+impl SQLOptions {
+ /// Create a new `SQLOptions` with default values
+ pub fn new() -> Self {
+ Default::default()
+ }
+
+ /// Should DML data modification commands (e.g. `INSERT and COPY`) be
run? Defaults to `true`.
+ pub fn with_allow_ddl(mut self, allow: bool) -> Self {
+ self.allow_ddl = allow;
+ self
+ }
+
+ /// Should DML data modification commands (e.g. `INSERT and COPY`) be run?
Defaults to `true`
+ pub fn with_allow_dml(mut self, allow: bool) -> Self {
+ self.allow_dml = allow;
+ self
+ }
+
+ /// Should Statements such as (e.g. `SET VARIABLE and `BEGIN TRANSACTION`
...`) be run?. Defaults to `true`
+ pub fn with_allow_statements(mut self, allow: bool) -> Self {
+ self.allow_statements = allow;
+ self
+ }
+
+ /// Return an error if the [`LogicalPlan`] has any nodes that are
+ /// incompatible with this [`SQLOptions`].
+ pub fn verify_plan(&self, plan: &LogicalPlan) -> Result<()> {
+ plan.visit(&mut BadPlanVisitor::new(self))?;
+ Ok(())
+ }
+}
+
+struct BadPlanVisitor<'a> {
+ options: &'a SQLOptions,
+}
+impl<'a> BadPlanVisitor<'a> {
+ fn new(options: &'a SQLOptions) -> Self {
+ Self { options }
+ }
+}
+
+impl<'a> TreeNodeVisitor for BadPlanVisitor<'a> {
+ type N = LogicalPlan;
+
+ fn pre_visit(&mut self, node: &Self::N) -> Result<VisitRecursion> {
+ match node {
+ LogicalPlan::Ddl(ddl) if !self.options.allow_ddl => {
+ plan_err!("DDL not supported: {}", ddl.name())
+ }
+ LogicalPlan::Dml(dml) if !self.options.allow_dml => {
+ plan_err!("DML not supported: {}", dml.op)
+ }
+ LogicalPlan::Copy(_) if !self.options.allow_dml => {
Review Comment:
Thanks so much @alamb for adding this to the core DF APIs. The addition of
the Copy operator is a great example of why it belongs here: we already have a
visitor that does similar filtering, but we would probably have failed to
update it to include Copy when we migrate to the latest DF version.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
