add a small test

Author: tonyalaribe

datafusion-postgres/src/handlers.rs

@@ -1009,4 +1009,46 @@ mod tests {
         let timeout = DfSessionService::get_statement_timeout(&client);
         assert_eq!(timeout, None);
     }
+
+    struct TestHook;
+
+    #[async_trait]
+    impl QueryHook for TestHook {
+        async fn handle_query(
+            &self,
+            statement: &Statement,
+            _ctx: &SessionContext,
+            _client: &dyn ClientInfo,
+        ) -> Option<PgWireResult<Vec<Response>>> {
+            if statement.to_string().contains("magic") {
+                Some(Ok(vec![Response::EmptyQuery]))
+            } else {
+                None
+            }
+        }
+    }
+
+    #[tokio::test]
+    async fn test_query_hooks() {
+        let hook = TestHook;
+        let ctx = SessionContext::new();
+        let client = MockClient::new();
+        
+        // Parse a statement that contains "magic"
+        let parser = PostgresCompatibilityParser::new();
+        let statements = parser.parse("SELECT magic").unwrap();
+        let stmt = Statement::Statement(Box::new(statements[0].clone()));
+        
+        // Hook should intercept
+        let result = hook.handle_query(&stmt, &ctx, &client).await;
+        assert!(result.is_some());
+        
+        // Parse a normal statement
+        let statements = parser.parse("SELECT 1").unwrap();
+        let stmt = Statement::Statement(Box::new(statements[0].clone()));
+        
+        // Hook should not intercept
+        let result = hook.handle_query(&stmt, &ctx, &client).await;
+        assert!(result.is_none());
+    }
 }