add pre_execute and post_execute hooks to plugins

buremba · Feb 7, 2025 · 7f0cc12 · 7f0cc12
1 parent 8e3af2f
commit 7f0cc12
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 0 deletions.
diff --git a/universql/plugin.py b/universql/plugin.py
@@ -94,6 +94,12 @@ def __init__(self,
     def transform_sql(self, expression: Expression, target_executor: Executor) -> Expression:
         return expression
 
+    def post_execute(self, expression: Expression, locations : typing.Optional[Locations], target_executor: Executor):
+        pass
+
+    def pre_execute(self, expression: Expression, target_executor: Executor):
+        pass
+
 
 # {"duckdb": DuckdbCatalog ..}
 COMPUTES = {}

diff --git a/universql/protocol/session.py b/universql/protocol/session.py
@@ -188,7 +188,23 @@ def perform_query(self, alternative_executor: Executor, raw_query, ast=None) ->
                             message = f"Unable to perform transformation {transform.__class__}"
                             logger.error(message, exc_info=e)
                             raise QueryError(f"{message}: {str(e)}")
+                    for transform in self.transforms:
+                        try:
+                            current_ast = transform.pre_execute(current_ast, alternative_executor)
+                        except Exception as e:
+                            print_exc(10)
+                            message = f"Unable to perform transformation {transform.__class__}"
+                            logger.error(message, exc_info=e)
+                            raise QueryError(f"{message}: {str(e)}")
                     new_locations = alternative_executor.execute(current_ast, self.catalog_executor, locations)
+                    for transform in self.transforms:
+                        try:
+                            current_ast = transform.post_execute(current_ast, new_locations, alternative_executor)
+                        except Exception as e:
+                            print_exc(10)
+                            message = f"Unable to perform transformation {transform.__class__}"
+                            logger.error(message, exc_info=e)
+                            raise QueryError(f"{message}: {str(e)}")
                 if new_locations is not None:
                     with sentry_sdk.start_span(op=op_name, name="Register new locations"):
                         self.catalog.register_locations(new_locations)