clean up syntax

Author: alamb

datafusion/core/src/execution/context.rs

@@ -69,7 +69,11 @@ use crate::logical_expr::{
     LogicalPlanBuilder, SetVariable, TableSource, TableType, UNNAMED_TABLE,
 };
 use crate::optimizer::OptimizerRule;
-use datafusion_sql::{planner::ParserOptions, ResolvedTableReference, TableReference};
+use datafusion_sql::{
+    parser::{CopyToSource, CopyToStatement},
+    planner::ParserOptions,
+    ResolvedTableReference, TableReference,
+};
 
 use crate::physical_optimizer::coalesce_batches::CoalesceBatches;
 use crate::physical_optimizer::repartition::Repartition;
@@ -1649,45 +1653,58 @@ impl SessionState {
         // table providers for all relations referenced in this query
         let mut relations = hashbrown::HashSet::with_capacity(10);
 
-        match statement {
-            DFStatement::Statement(s) => {
-                struct RelationVisitor<'a>(&'a mut hashbrown::HashSet<ObjectName>);
+        struct RelationVisitor<'a>(&'a mut hashbrown::HashSet<ObjectName>);
+
+        impl<'a> RelationVisitor<'a> {
+            /// Record that `relation` was used in this statement
+            fn insert(&mut self, relation: &ObjectName) {
+                self.0.get_or_insert_with(relation, |_| relation.clone());
+            }
+        }
 
-                impl<'a> Visitor for RelationVisitor<'a> {
-                    type Break = ();
+        impl<'a> Visitor for RelationVisitor<'a> {
+            type Break = ();
 
-                    fn pre_visit_relation(
-                        &mut self,
-                        relation: &ObjectName,
-                    ) -> ControlFlow<()> {
-                        self.0.get_or_insert_with(relation, |_| relation.clone());
-                        ControlFlow::Continue(())
-                    }
+            fn pre_visit_relation(&mut self, relation: &ObjectName) -> ControlFlow<()> {
+                self.insert(relation);
+                ControlFlow::Continue(())
+            }
 
-                    fn pre_visit_statement(
-                        &mut self,
-                        statement: &Statement,
-                    ) -> ControlFlow<()> {
-                        if let Statement::ShowCreate {
-                            obj_type: ShowCreateObject::Table | ShowCreateObject::View,
-                            obj_name,
-                        } = statement
-                        {
-                            self.0.get_or_insert_with(obj_name, |_| obj_name.clone());
-                        }
-                        ControlFlow::Continue(())
-                    }
+            fn pre_visit_statement(&mut self, statement: &Statement) -> ControlFlow<()> {
+                if let Statement::ShowCreate {
+                    obj_type: ShowCreateObject::Table | ShowCreateObject::View,
+                    obj_name,
+                } = statement
+                {
+                    self.insert(obj_name)
                 }
-                let mut visitor = RelationVisitor(&mut relations);
+                ControlFlow::Continue(())
+            }
+        }
+
+        let mut visitor = RelationVisitor(&mut relations);
+        match statement {
+            DFStatement::Statement(s) => {
                 let _ = s.as_ref().visit(&mut visitor);
             }
             DFStatement::CreateExternalTable(table) => {
-                relations.insert(ObjectName(vec![Ident::from(table.name.as_str())]));
-            }
-            DFStatement::DescribeTableStmt(table) => {
-                relations
-                    .get_or_insert_with(&table.table_name, |_| table.table_name.clone());
+                visitor
+                    .0
+                    .insert(ObjectName(vec![Ident::from(table.name.as_str())]));
             }
+            DFStatement::DescribeTableStmt(table) => visitor.insert(&table.table_name),
+            DFStatement::CopyTo(CopyToStatement {
+                source,
+                target: _,
+                options: _,
+            }) => match source {
+                CopyToSource::Relation(table_name) => {
+                    visitor.insert(table_name);
+                }
+                CopyToSource::Query(query) => {
+                    query.visit(&mut visitor);
+                }
+            },
         }
 
         // Always include information_schema if available

datafusion/core/tests/sqllogictests/test_files/copy.slt

@@ -20,61 +20,21 @@
 statement ok
 create table source_table(col1 integer, col2 varchar) as values (1, 'Foo'), (2, 'Bar');
 
-# create an external parquet file
-# TODO figure out how to use a real tempfile name
-#statement ok
-#COPY source_table to '/tmp/table.parquet';
+# Copy from table
+statement error DataFusion error: This feature is not implemented: `COPY \.\. TO \.\.` statement is yet supported
+COPY source_table  to '/tmp/table.parquet';
 
+# Copy from table with options
+statement error DataFusion error: This feature is not implemented: `COPY \.\. TO \.\.` statement is yet supported
+COPY source_table  to '/tmp/table.parquet' (row_group_size 55);
 
-# create an external parquet file from a query
-# TODO figure out how to use a real tempfile name
-statement ok
-COPY (select col2, sum(col1) from source_table group by col2 order by col2)  to '/tmp/table.parquet';
-
-# Read the data back in
-statement ok
-CREATE external table new_table STORED as PARQUET LOCATION '/tmp/table.parquet';
-
-# Read the data back in
-query IT
-select * from new_table;
-----
-1 Foo
-2 Bar
-
-statement ok
-drop table new_table;
-
-
-# TODO error cases:
-# incomplete statement
-# pass non literal to COPY options
-
-
-
-
-
-# TODO add support for reading directly for files without needing to create an external table
-# (TODO find ticket)
 
+# Error cases:
 
+# Incomplete statement
+statement error DataFusion error: SQL error: ParserError\("Expected \), found: EOF"\)
+COPY (select col2, sum(col1) from source_table
 
-# TODO test copying to partitioned target (directory)
-
-
-# TODO test copying with options
-
-
-
-# TODO test query source
-
-
-
-# TODO: support CSV, JSON and AVRO
-
-# TODO: document different options
-
-
-
-statement ok
-drop table source_table;
+# Copy from table with non literal
+statement error DataFusion error: SQL error: ParserError\("Expected ',' or '\)' after option definition, found: \+"\)
+COPY source_table  to '/tmp/table.parquet' (row_group_size 55 + 102);