Repository: inspektor-dev/inspektor (language: rust)

Issue:
add transaction support in postgres driver
Current postgres driver, only checks incoming query is valid or not. But, it doesn't know how to act in a transaction block.

for e.g.: end user can query user but can't insert as policy.

```
SELECT * FROM USERS;
BEGIN;
INSERT INTO USERS(first_name, last_name) VALUES ('pooni', 'kuttypoonai');
COMMIT;
```

The expected behaviour is that postgres should return data for select statements and send error for insert transaction. But, Inspektor 
Postgres driver return error response directly.


Discussion / hints:

Fix patch:
diff --git a/.gitignore b/.gitignore
index dc58b8e..8176aaf 100644
--- a/.gitignore
+++ b/.gitignore
@@ -2,4 +2,5 @@
 /controlplane/config.yaml
 /controlplane/inspektor
 /controlplane/dashboard/node_modules
-/controlplane/policy_dir
\ No newline at end of file
+/controlplane/policy_dir
+dataplane_config.yaml
\ No newline at end of file
diff --git a/controlplane/dashboard/src/store/store.js b/controlplane/dashboard/src/store/store.js
index 00356a3..38311c4 100644
--- a/controlplane/dashboard/src/store/store.js
+++ b/controlplane/dashboard/src/store/store.js
@@ -70,7 +70,7 @@ const store = createStore({
             // merge session and meta in same object.
             for (let i = 0; i < datasources.length; i++) {
                 for (let j = 0; j < sessions.length; j++) {
-                    if (sessions[j].objectID == datasources[i].id && sessions[j].meta.expiresAt == 0) {
+                    if (sessions[j].objectID == datasources[i].id && (sessions[j].meta.expiresAt == 0 || sessions[j].meta.expiresAt == undefined)) {
                         datasources[i].sessionMeta = sessions[j].meta
                     }
                 }
@@ -78,7 +78,7 @@ const store = createStore({
             commit("setDatasource", datasources)
             let tempSesions = await api.getTempCredentials()
             let tempDatasources = []
-            for (let k= 0; k< tempSesions.length; k++){
+            for (let k = 0; k < tempSesions.length; k++) {
                 let datasource = tempSesions[k].datasource;
                 datasource.sessionMeta = tempSesions[k]
                 tempDatasources.push(datasource)
diff --git a/src/apiproto/api_grpc.rs b/src/apiproto/api_grpc.rs
index dcd9743..2835a97 100644
--- a/src/apiproto/api_grpc.rs
+++ b/src/apiproto/api_grpc.rs
@@ -4,7 +4,6 @@
 // https://github.com/Manishearth/rust-clippy/issues/702
 #![allow(unknown_lints)]
 #![allow(clippy::all)]
-
 #![allow(box_pointers)]
 #![allow(dead_code)]
 #![allow(missing_docs)]
@@ -16,25 +15,48 @@
 #![allow(unused_imports)]
 #![allow(unused_results)]
 
-const METHOD_INSPEKTOR_AUTH: ::grpcio::Method<super::api::AuthRequest, super::api::AuthResponse> = ::grpcio::Method {
-    ty: ::grpcio::MethodType::Unary,
-    name: "/api.Inspektor/Auth",
-    req_mar: ::grpcio::Marshaller { ser: ::grpcio::pb_ser, de: ::grpcio::pb_de },
-    resp_mar: ::grpcio::Marshaller { ser: ::grpcio::pb_ser, de: ::grpcio::pb_de },
-};
-
-const METHOD_INSPEKTOR_POLICY: ::grpcio::Method<super::api::Empty, super::api::InspektorPolicy> = ::grpcio::Method {
-    ty: ::grpcio::MethodType::ServerStreaming,
-    name: "/api.Inspektor/Policy",
-    req_mar: ::grpcio::Marshaller { ser: ::grpcio::pb_ser, de: ::grpcio::pb_de },
-    resp_mar: ::grpcio::Marshaller { ser: ::grpcio::pb_ser, de: ::grpcio::pb_de },
-};
-
-const METHOD_INSPEKTOR_GET_DATA_SOURCE: ::grpcio::Method<super::api::Empty, super::api::DataSourceResponse> = ::grpcio::Method {
+const METHOD_INSPEKTOR_AUTH: ::grpcio::Method<super::api::AuthRequest, super::api::AuthResponse> =
+    ::grpcio::Method {
+        ty: ::grpcio::MethodType::Unary,
+        name: "/api.Inspektor/Auth",
+        req_mar: ::grpcio::Marshaller {
+            ser: ::grpcio::pb_ser,
+            de: ::grpcio::pb_de,
+        },
+        resp_mar: ::grpcio::Marshaller {
+            ser: ::grpcio::pb_ser,
+            de: ::grpcio::pb_de,
+        },
+    };
+
+const METHOD_INSPEKTOR_POLICY: ::grpcio::Method<super::api::Empty, super::api::InspektorPolicy> =
+    ::grpcio::Method {
+        ty: ::grpcio::MethodType::ServerStreaming,
+        name: "/api.Inspektor/Policy",
+        req_mar: ::grpcio::Marshaller {
+            ser: ::grpcio::pb_ser,
+            de: ::grpcio::pb_de,
+        },
+        resp_mar: ::grpcio::Marshaller {
+            ser: ::grpcio::pb_ser,
+            de: ::grpcio::pb_de,
+        },
+    };
+
+const METHOD_INSPEKTOR_GET_DATA_SOURCE: ::grpcio::Method<
+    super::api::Empty,
+    super::api::DataSourceResponse,
+> = ::grpcio::Method {
     ty: ::grpcio::MethodType::Unary,
     name: "/api.Inspektor/GetDataSource",
-    req_mar: ::grpcio::Marshaller { ser: ::grpcio::pb_ser, de: ::grpcio::pb_de },
-    resp_mar: ::grpcio::Marshaller { ser: ::grpcio::pb_ser, de: ::grpcio::pb_de },
+    req_mar: ::grpcio::Marshaller {
+        ser: ::grpcio::pb_ser,
+        de: ::grpcio::pb_de,
+    },
+    resp_mar: ::grpcio::Marshaller {
+        ser: ::grpcio::pb_ser,
+        de: ::grpcio::pb_de,
+    },
 };
 
 #[derive(Clone)]
@@ -49,54 +71,111 @@ impl InspektorClient {
         }
     }
 
-    pub fn auth_opt(&self, req: &super::api::AuthRequest, opt: ::grpcio::CallOption) -> ::grpcio::Result<super::api::AuthResponse> {
+    pub fn auth_opt(
+        &self,
+        req: &super::api::AuthRequest,
+        opt: ::grpcio::CallOption,
+    ) -> ::grpcio::Result<super::api::AuthResponse> {
         self.client.unary_call(&METHOD_INSPEKTOR_AUTH, req, opt)
     }
 
-    pub fn auth(&self, req: &super::api::AuthRequest) -> ::grpcio::Result<super::api::AuthResponse> {
+    pub fn auth(
+        &self,
+        req: &super::api::AuthRequest,
+    ) -> ::grpcio::Result<super::api::AuthResponse> {
         self.auth_opt(req, ::grpcio::CallOption::default())
     }
 
-    pub fn auth_async_opt(&self, req: &super::api::AuthRequest, opt: ::grpcio::CallOption) -> ::grpcio::Result<::grpcio::ClientUnaryReceiver<super::api::AuthResponse>> {
-        self.client.unary_call_async(&METHOD_INSPEKTOR_AUTH, req, opt)
+    pub fn auth_async_opt(
+        &self,
+        req: &super::api::AuthRequest,
+        opt: ::grpcio::CallOption,
+    ) -> ::grpcio::Result<::grpcio::ClientUnaryReceiver<super::api::AuthResponse>> {
+        self.client
+            .unary_call_async(&METHOD_INSPEKTOR_AUTH, req, opt)
     }
 
-    pub fn auth_async(&self, req: &super::api::AuthRequest) -> ::grpcio::Result<::grpcio::ClientUnaryReceiver<super::api::AuthResponse>> {
+    pub fn auth_async(
+        &self,
+        req: &super::api::AuthRequest,
+    ) -> ::grpcio::Result<::grpcio::ClientUnaryReceiver<super::api::AuthResponse>> {
         self.auth_async_opt(req, ::grpcio::CallOption::default())
     }
 
-    pub fn policy_opt(&self, req: &super::api::Empty, opt: ::grpcio::CallOption) -> ::grpcio::Result<::grpcio::ClientSStreamReceiver<super::api::InspektorPolicy>> {
-        self.client.server_streaming(&METHOD_INSPEKTOR_POLICY, req, opt)
+    pub fn policy_opt(
+        &self,
+        req: &super::api::Empty,
+        opt: ::grpcio::CallOption,
+    ) -> ::grpcio::Result<::grpcio::ClientSStreamReceiver<super::api::InspektorPolicy>> {
+        self.client
+            .server_streaming(&METHOD_INSPEKTOR_POLICY, req, opt)
     }
 
-    pub fn policy(&self, req: &super::api::Empty) -> ::grpcio::Result<::grpcio::ClientSStreamReceiver<super::api::InspektorPolicy>> {
+    pub fn policy(
+        &self,
+        req: &super::api::Empty,
+    ) -> ::grpcio::Result<::grpcio::ClientSStreamReceiver<super::api::InspektorPolicy>> {
         self.policy_opt(req, ::grpcio::CallOption::default())
     }
 
-    pub fn get_data_source_opt(&self, req: &super::api::Empty, opt: ::grpcio::CallOption) -> ::grpcio::Result<super::api::DataSourceResponse> {
-        self.client.unary_call(&METHOD_INSPEKTOR_GET_DATA_SOURCE, req, opt)
+    pub fn get_data_source_opt(
+        &self,
+        req: &super::api::Empty,
+        opt: ::grpcio::CallOption,
+    ) -> ::grpcio::Result<super::api::DataSourceResponse> {
+        self.client
+            .unary_call(&METHOD_INSPEKTOR_GET_DATA_SOURCE, req, opt)
     }
 
-    pub fn get_data_source(&self, req: &super::api::Empty) -> ::grpcio::Result<super::api::DataSourceResponse> {
+    pub fn get_data_source(
+        &self,
+        req: &super::api::Empty,
+    ) -> ::grpcio::Result<super::api::DataSourceResponse> {
         self.get_data_source_opt(req, ::grpcio::CallOption::default())
     }
 
-    pub fn get_data_source_async_opt(&self, req: &super::api::Empty, opt: ::grpcio::CallOption) -> ::grpcio::Result<::grpcio::ClientUnaryReceiver<super::api::DataSourceResponse>> {
-        self.client.unary_call_async(&METHOD_INSPEKTOR_GET_DATA_SOURCE, req, opt)
+    pub fn get_data_source_async_opt(
+        &self,
+        req: &super::api::Empty,
+        opt: ::grpcio::CallOption,
+    ) -> ::grpcio::Result<::grpcio::ClientUnaryReceiver<super::api::DataSourceResponse>> {
+        self.client
+            .unary_call_async(&METHOD_INSPEKTOR_GET_DATA_SOURCE, req, opt)
     }
 
-    pub fn get_data_source_async(&self, req: &super::api::Empty) -> ::grpcio::Result<::grpcio::ClientUnaryReceiver<super::api::DataSourceResponse>> {
+    pub fn get_data_source_async(
+        &self,
+        req: &super::api::Empty,
+    ) -> ::grpcio::Result<::grpcio::ClientUnaryReceiver<super::api::DataSourceResponse>> {
         self.get_data_source_async_opt(req, ::grpcio::CallOption::default())
     }
-    pub fn spawn<F>(&self, f: F) where F: ::futures::Future<Output = ()> + Send + 'static {
+    pub fn spawn<F>(&self, f: F)
+    where
+        F: ::futures::Future<Output = ()> + Send + 'static,
+    {
         self.client.spawn(f)
     }
 }
 
 pub trait Inspektor {
-    fn auth(&mut self, ctx: ::grpcio::RpcContext, req: super::api::AuthRequest, sink: ::grpcio::UnarySink<super::api::AuthResponse>);
-    fn policy(&mut self, ctx: ::grpcio::RpcContext, req: super::api::Empty, sink: ::grpcio::ServerStreamingSink<super::api::InspektorPolicy>);
-    fn get_data_source(&mut self, ctx: ::grpcio::RpcContext, req: super::api::Empty, sink: ::grpcio::UnarySink<super::api::DataSourceResponse>);
+    fn auth(
+        &mut self,
+        ctx: ::grpcio::RpcContext,
+        req: super::api::AuthRequest,
+        sink: ::grpcio::UnarySink<super::api::AuthResponse>,
+    );
+    fn policy(
+        &mut self,
+        ctx: ::grpcio::RpcContext,
+        req: super::api::Empty,
+        sink: ::grpcio::ServerStreamingSink<super::api::InspektorPolicy>,
+    );
+    fn get_data_source(
+        &mut self,
+        ctx: ::grpcio::RpcContext,
+        req: super::api::Empty,
+        sink: ::grpcio::UnarySink<super::api::DataSourceResponse>,
+    );
 }
 
 pub fn create_inspektor<S: Inspektor + Send + Clone + 'static>(s: S) -> ::grpcio::Service {
@@ -106,12 +185,14 @@ pub fn create_inspektor<S: Inspektor + Send + Clone + 'static>(s: S) -> ::grpcio
         instance.auth(ctx, req, resp)
     });
     let mut instance = s.clone();
-    builder = builder.add_server_streaming_handler(&METHOD_INSPEKTOR_POLICY, move |ctx, req, resp| {
-        instance.policy(ctx, req, resp)
-    });
+    builder = builder
+        .add_server_streaming_handler(&METHOD_INSPEKTOR_POLICY, move |ctx, req, resp| {
+            instance.policy(ctx, req, resp)
+        });
     let mut instance = s;
-    builder = builder.add_unary_handler(&METHOD_INSPEKTOR_GET_DATA_SOURCE, move |ctx, req, resp| {
-        instance.get_data_source(ctx, req, resp)
-    });
+    builder = builder
+        .add_unary_handler(&METHOD_INSPEKTOR_GET_DATA_SOURCE, move |ctx, req, resp| {
+            instance.get_data_source(ctx, req, resp)
+        });
     builder.build()
 }
diff --git a/src/postgres_driver/message.rs b/src/postgres_driver/message.rs
index eca51f6..5ebb125 100644
--- a/src/postgres_driver/message.rs
+++ b/src/postgres_driver/message.rs
@@ -34,6 +34,15 @@ impl ReadyState {
             ReadyState::FailedTransaction => b'E',
         }
     }
+
+    fn from_u8(state: u8) -> ReadyState {
+        match state {
+            b'I' => ReadyState::Idle,
+            b'T' => ReadyState::Transaction,
+            b'E' => ReadyState::FailedTransaction,
+            _ => unreachable!("invalid ready state"),
+        }
+    }
 }
 
 #[derive(Debug)]
@@ -46,6 +55,7 @@ pub enum BackendMessage {
     AuthenticationSASLContinue { data: Vec<u8> },
     AuthenticationSASLFinal { data: Vec<u8> },
     ReadyForQuery { state: ReadyState },
+    Message { data: Vec<u8>, meta: u8 },
 }
 
 impl BackendMessage {
@@ -91,6 +101,16 @@ impl BackendMessage {
                 buf.put_u8(state.get_state_byte());
                 buf
             }
+            BackendMessage::Message { meta, data } => {
+                debug!("encoding emeta {:?}", *meta as char);
+                buf.put_u8(*meta);
+                write_message(&mut buf, |buf| {
+                    buf.extend_from_slice(&data[..]);
+                    Ok(())
+                })
+                .unwrap();
+                buf
+            }
             _ => {
                 unreachable!("encoding invalid startup message")
             }
@@ -108,20 +128,18 @@ where
         error!("error while reading the meta for backend message {:?}", e);
         anyhow!("error reading backend meta")
     })?;
-
+    let len = decode_frame_length(&mut conn).await.map_err(|e| {
+        error!("error while decoding frame lenght {:?}", e);
+        anyhow!("invalid backend message")
+    })?;
+    let mut buf = BytesMut::new();
+    buf.resize(len, b'0');
+    conn.read_exact(&mut buf).await.map_err(|err| {
+        error!("error while reading backend message {:?}", err);
+        anyhow!("error while decoding error messaage")
+    })?;
     match meta[0] {
         b'R' => {
-            let len = decode_frame_length(&mut conn).await.map_err(|e| {
-                error!("error while decoding frame lenght {:?}", e);
-                anyhow!("invalid backend message")
-            })?;
-            let mut buf = BytesMut::new();
-            buf.resize(len, b'0');
-            conn.read_exact(&mut buf).await.map_err(|err| {
-                error!("error while reading backend message {:?}", err);
-                anyhow!("error while decoding error messaage")
-            })?;
-
             let msg_type = NetworkEndian::read_u32(&buf);
             buf.advance(4);
             match msg_type {
@@ -153,15 +171,6 @@ where
             }
         }
         b'E' => {
-            let msg_len = decode_frame_length(&mut conn).await.map_err(|e| {
-                error!("error while decoding backend error message length");
-                e
-            })?;
-            let mut buf = BytesMut::new();
-            buf.resize(msg_len, b'0');
-            conn.read_exact(&mut buf)
-                .await
-                .map_err(|_| anyhow!("error while decoing err msg from postgres target"))?;
             if buf[0] == 0 {
                 return Ok(BackendMessage::ErrorMsg(None));
             }
@@ -169,9 +178,15 @@ where
             let err_msg = read_cstr(&mut buf)?;
             return Ok(BackendMessage::ErrorMsg(Some(err_msg)));
         }
+        b'Z' => {
+            let state = ReadyState::from_u8(buf[0]);
+            return Ok(BackendMessage::ReadyForQuery { state: state });
+        }
         _ => {
-            error!("invalid meta message for backend {:?}", meta[0]);
-            return Err(anyhow!("invalid backend message"));
+            return Ok(BackendMessage::Message {
+                data: buf.to_vec(),
+                meta: meta[0],
+            });
         }
     }
 }
diff --git a/src/postgres_driver/protocol_handler.rs b/src/postgres_driver/protocol_handler.rs
index d3d90e4..e4e4b72 100644
--- a/src/postgres_driver/protocol_handler.rs
+++ b/src/postgres_driver/protocol_handler.rs
@@ -24,6 +24,7 @@ use log::*;
 use md5::{Digest, Md5};
 use openssl::ssl::{Ssl, SslConnector, SslMethod};
 use postgres_protocol::authentication::sasl;
+use sqlparser::ast::Statement;
 use std::collections::{HashMap, HashSet};
 use std::pin::Pin;
 use std::time::Duration;
@@ -52,6 +53,7 @@ pub struct ProtocolHandler {
     config: PostgresConfig,
     connected_db: String,
     datasource_name: String,
+    pending_error: Option<ProtocolHandlerError>,
 }
 
 #[derive(Default)]
@@ -200,20 +202,35 @@ impl ProtocolHandler {
                 }
                     self.policy_evaluator = evaluator;
                 }
-                n = self.target_conn.read(&mut target_buf) => {
+                n = decode_backend_message(&mut self.target_conn) => {
                     match n {
                         Err(e) =>{
                                 println!("failed to read from socket; err = {:?}", e);
                                 return Ok(());
                         },
-                        Ok(n) =>{
-                            if n == 0 {
-                               return Ok(());
+                        Ok(msg) =>{
+                            if self.pending_error.is_some(){
+                                // check the incoming message is ready for query.
+                                // if it's ready for query send the error message before
+                                // forwarding ready for query message.
+                                match &msg{
+                                    BackendMessage::ReadyForQuery{..} => {
+                                        // send the pending error message.
+                                        let e = self.pending_error.as_ref().unwrap();
+                                        let err_rsp = BackendMessage::ErrorMsg(Some(format!("{}", e)));
+                                        if let Err(e) = self.client_conn.write_all(&err_rsp.encode()).await{
+                                            error!("error while writing the rsp message to the client {:?}", e);
+                                            return Ok(());
+                                        }
+                                    },
+                                    _=> {}
+                                }
                             }
-                            if let Err(e) = self.client_conn.write_all(&target_buf[0..n]).await{
+                            if let Err(e) = self.client_conn.write_all(&msg.encode()).await{
                                 error!("error while writing the rsp message to the client {:?}", e);
                                 return Ok(());
                             }
+                            continue;
                         }
                     }
                 }
@@ -383,6 +400,7 @@ impl ProtocolHandler {
                         config: config.clone(),
                         connected_db: client_parms.get("database").unwrap().clone(),
                         datasource_name: datasource_name,
+                        pending_error: None,
                     };
                     return Ok(handler);
                 }
@@ -552,10 +570,6 @@ impl ProtocolHandler {
         schemas: Vec<String>,
     ) -> Result<(), ProtocolHandlerError> {
         debug!("input query {}", query);
-        // TODO: this needs to worked out.
-        if query.contains("BEGIN READ ONLY") {
-            return Ok(());
-        }
         let dialect = sqlparser::dialect::PostgreSqlDialect {};
         let mut statements = match sqlparser::parser::Parser::parse_sql(&dialect, query) {
             Ok(statements) => statements,
@@ -570,9 +584,17 @@ impl ProtocolHandler {
         let rule = self.get_rule_engine()?;
         debug!("rewriting with schema {:?}", schemas);
         let rewriter = QueryRewriter::new(rule, schemas);
-        rewriter.rewrite(&mut statements, ctx)?;
         let mut out = String::from("");
-        for statement in statements {
+        let mut good_to_forward = false;
+        for statement in &mut statements {
+            if let Err(e) = rewriter.rewrite(statement, &ctx) {
+                if !good_to_forward {
+                    return Err(ProtocolHandlerError::RewriterError(e));
+                }
+                self.pending_error = Some(ProtocolHandlerError::RewriterError(e));
+                break;
+            }
+            good_to_forward = true;
             out = format!("{}{};", out, statement);
         }
         debug!("output query {}", out);
diff --git a/src/sql/query_rewriter.rs b/src/sql/query_rewriter.rs
index 1f386bd..18f8928 100644
--- a/src/sql/query_rewriter.rs
+++ b/src/sql/query_rewriter.rs
@@ -42,52 +42,48 @@ impl<T: RuleEngine + Clone> QueryRewriter<T> {
 
     pub fn rewrite(
         &self,
-        statements: &mut Vec<Statement>,
-        state: Ctx,
+        statement: &mut Statement,
+        state: &Ctx,
     ) -> Result<(), QueryRewriterError> {
-        for statement in statements {
-            match statement {
-                Statement::Query(query) => {
-                    self.handle_query(query, &state)?;
+        match statement {
+            Statement::Query(query) => {
+                self.handle_query(query, state)?;
+            }
+            Statement::Update {
+                table, assignments, ..
+            } => {
+                if !self.rule_engine.is_update_allowed() {
+                    return Err(QueryRewriterError::UnAuthorizedUpdate);
                 }
-                Statement::Update {
-                    table, assignments, ..
-                } => {
-                    if !self.rule_engine.is_update_allowed() {
-                        return Err(QueryRewriterError::UnAuthorizedUpdate);
-                    }
-                    self.handle_update(&table, &assignments)?;
+                self.handle_update(&table, &assignments)?;
+            }
+            Statement::Insert {
+                columns,
+                table_name,
+                ..
+            } => {
+                if !self.rule_engine.is_insert_allowed() {
+                    return Err(QueryRewriterError::UnAuthorizedInsert);
                 }
-                Statement::Insert {
-                    columns,
-                    table_name,
-                    ..
-                } => {
-                    if !self.rule_engine.is_insert_allowed() {
-                        return Err(QueryRewriterError::UnAuthorizedInsert);
-                    }
-                    let allowed_attributes = self.rule_engine.get_allowed_insert_attributes();
-                    if !self.is_operation_allowed(&table_name, &columns, allowed_attributes) {
-                        return Err(QueryRewriterError::UnAuthorizedInsert);
-                    }
+                let allowed_attributes = self.rule_engine.get_allowed_insert_attributes();
+                if !self.is_operation_allowed(&table_name, &columns, allowed_attributes) {
+                    return Err(QueryRewriterError::UnAuthorizedInsert);
                 }
-                Statement::Copy {
-                    table_name,
-                    columns,
-                    ..
-                } => {
-                    if !self.rule_engine.is_copy_allowed() {
-                        return Err(QueryRewriterError::UnAthorizedCopy);
-                    }
-                    let allowed_attributes = self.rule_engine.get_allowed_copy_attributes();
-                    if !self.is_operation_allowed(&table_name, &columns, allowed_attributes) {
-                        return Err(QueryRewriterError::UnAthorizedCopy);
-                    }
+            }
+            Statement::Copy {
+                table_name,
+                columns,
+                ..
+            } => {
+                if !self.rule_engine.is_copy_allowed() {
+                    return Err(QueryRewriterError::UnAthorizedCopy);
                 }
-                _ => {
-                    continue;
+                let allowed_attributes = self.rule_engine.get_allowed_copy_attributes();
+                if !self.is_operation_allowed(&table_name, &columns, allowed_attributes) {
+                    return Err(QueryRewriterError::UnAthorizedCopy);
                 }
             }
+            _ => {}
         }
         Ok(())
     }
@@ -673,7 +669,9 @@ mod tests {
     ) {
         let dialect = PostgreSqlDialect {};
         let mut statements = Parser::parse_sql(&dialect, input).unwrap();
-        rewriter.rewrite(&mut statements, state).unwrap();
+        for statement in &mut statements {
+            rewriter.rewrite(statement, &state).unwrap();
+        }
         assert_eq!(output, format!("{}", statements[0]))
     }
 
@@ -685,8 +683,10 @@ mod tests {
     ) {
         let dialect = PostgreSqlDialect {};
         let mut statements = Parser::parse_sql(&dialect, input).unwrap();
-        let rewriter_err = rewriter.rewrite(&mut statements, state).unwrap_err();
-        assert_eq!(rewriter_err, err)
+        for statement in &mut statements {
+            let rewriter_err = rewriter.rewrite(statement, &state).unwrap_err();
+            assert_eq!(rewriter_err, err)
+        }
     }
 
     #[test]


Test patch: