vdk-core: add new managed db_connection_execute_operation hook

Currently we cannot track the full execution of a query (before and
after it) . This is necessary in order to be able to generate proper
lineage events (query start and end) or to track its duration (we had to
currently had to add this to vdk-core
(#804) but that's
really better implemented as plugin (think aspect-oriented programming).

In the future we should consider adding hook for fetch (fetchMany,
fetchAll) as well as it's there where the data is returned (plugins can
take stats - count number of rows, validate sensitive columns, do result
validation tests, etc.)

Testing Done: unit tests, functional test.

Signed-off-by: Antoni Ivanov <[email protected]>

Author: antoniivanov

projects/vdk-core/src/vdk/api/plugin/connection_hook_spec.py

@@ -5,6 +5,10 @@
 
 from vdk.api.plugin.hook_markers import hookspec
 from vdk.internal.builtin_plugins.connection.decoration_cursor import DecorationCursor
+from vdk.internal.builtin_plugins.connection.execution_cursor import (
+    ExecuteOperationResult,
+)
+from vdk.internal.builtin_plugins.connection.execution_cursor import ExecutionCursor
 from vdk.internal.builtin_plugins.connection.recovery_cursor import RecoveryCursor
 
 
@@ -68,6 +72,68 @@ def db_connection_decorate_operation(
         """
         pass
 
+    @hookspec(firstresult=True)
+    def db_connection_execute_operation(
+        self, execution_cursor: ExecutionCursor
+    ) -> Optional[ExecuteOperationResult]:
+        """
+        The method that executes the actual SQL query using execution cursor.
+
+        When the hook is called the call loop will only invoke up to
+        the first hookimpl which returns a result other then None.
+        It is then taken as result of the overall hook call.
+
+        You can see the default implementation at DefaultConnectionHookImpl.
+
+        If you want to override the default implementation return a non-None result using tryfirst=True or no decorator.
+        But in most cases you would not need to overwrite the default implementation (hence return None).
+        If some hookimpl function raises an exception, the behaviour is as outlined in hook_markers module
+
+        Often it may be needed to wrap around it so that you can track the execution of the query.
+
+        For example:
+
+        .. code-block::
+
+            @hookimpl(hookwrapper=True)
+            db_connection_execute_operation(execution_cursor: ExecutionCursor) -> Optional[int]:
+                # let's track duration of the query
+                start = time.time()
+                log.info(f"Starting query: {execution_cursor.get_managed_operation().get_operation()}")
+                outcome: pluggy.callers._Result
+                outcome = yield # we yield the execution to other implementations (including default one)
+                is_success: bool = outcome.excinfo is None
+                log.info(f"Query finished. duration: {time.time() - start}. is_success: {is_success}")
+                # no return!
+
+        Another example - we want to change the result
+
+        .. code-block::
+
+            @hookimpl(hookwrapper=True)
+            db_connection_execute_operation(execution_cursor: ExecutionCursor) -> Optional[int]:
+                outcome: pluggy.callers._Result
+                outcome = yield #
+                outcome.force_result(new_result)  # set new return result
+
+        Another example - let's say we are writing vdk-impala plugin and want to print more debug info
+        which is available from the Impala native cursor (provided by impyla library)
+
+        .. code-block::
+
+            @hookimpl(hookwrapper=True)
+            db_connection_execute_operation(execution_cursor: ExecutionCursor) -> Optional[int]:
+                yield # let the query execute first
+                c = cast(impala.interface.Cursor, execution_cursor)
+                log.info(f"Query {execution_cursor.get_managed_operation().get_operation()} debug info:"
+                         f"summary: {c.get_summary()}, profile: {c.get_profile()}")
+
+
+        :param execution_cursor: ExecutionCursor
+            A PEP249Cursor implementation purposed for actual query execution.
+        """
+        pass
+
     @hookspec
     def db_connection_recover_operation(self, recovery_cursor: RecoveryCursor) -> None:
         """

projects/vdk-core/src/vdk/api/plugin/hook_markers.py

@@ -27,6 +27,10 @@
 The method name must match one of those defined as hookspec
 The method signature - the arguments need to be subset of arguments defined in hookspec
 
+If any hookimpl errors with an exception no further callbacks are invoked
+and the exception is packaged up and delivered to any wrappers
+before being re-raised at the hook invocation point.
+
 Plugin execution order can be configured in the decorator with following variables:
 
 If ``optionalhook`` is ``True`` a missing matching hook specification will not result
@@ -55,5 +59,7 @@ def func():
     after_other_hooks_executed(outcome.get_result())
     outcome.force_result(new_res)  # set new return result
 
+For more details see https://pluggy.readthedocs.io
+
 """
 hookimpl = pluggy.HookimplMarker(GROUP_NAME)

projects/vdk-core/src/vdk/internal/builtin_plugins/connection/connection_hooks.py

@@ -0,0 +1,71 @@
+# Copyright 2021 VMware, Inc.
+# SPDX-License-Identifier: Apache-2.0
+from typing import Any
+from typing import cast
+
+from vdk.api.plugin.connection_hook_spec import ConnectionHookSpec
+from vdk.api.plugin.hook_markers import hookimpl
+from vdk.api.plugin.plugin_registry import IPluginRegistry
+from vdk.internal.builtin_plugins.connection.execution_cursor import (
+    ExecuteOperationResult,
+)
+from vdk.internal.builtin_plugins.connection.execution_cursor import ExecutionCursor
+from vdk.internal.core.errors import ErrorMessage
+from vdk.internal.core.errors import PlatformServiceError
+
+
+class DefaultConnectionHookImpl:
+    """
+    The default implementation of execute operation.
+    Generally it should not be overridden.
+    To "override" with new implementation, make a new hook which returns non-None ExecuteOperationResult.
+    See ConnectionHookSpec documentation for more details.
+    """
+
+    @hookimpl(trylast=True)
+    def db_connection_execute_operation(self, execution_cursor: ExecutionCursor) -> Any:
+        managed_operation = execution_cursor.get_managed_operation()
+        native_cursor = execution_cursor.get_native_cursor()
+        if managed_operation.get_parameters():
+            native_result = native_cursor.execute(
+                managed_operation.get_operation(), managed_operation.get_parameters()
+            )
+        else:
+            native_result = native_cursor.execute(managed_operation.get_operation())
+        return ExecuteOperationResult(native_result)
+
+
+class ConnectionHookSpecFactory:
+    """
+    Class used to create properly initialized ConnectionHookSpec instance to use to execute the underlying hooks
+    """
+
+    def __init__(self, plugin_registry: IPluginRegistry):
+        self.__plugin_registry = plugin_registry
+
+    def get_connection_hook_spec(self) -> ConnectionHookSpec:
+        """
+        Returns ConnectionHookSpec class which would act as a relay and invoke the underlying implemented hooks
+        It's generally a ConnectionHookSpec cast of plugin_registry.PluginHookRelay to enable easier hook invocations.
+        It also initializes some of the hooks (now only db_connection_execute_operation) with default implementations.
+        :return: ConnectionHookSpec
+        """
+        if self.__plugin_registry:
+            if not self.__plugin_registry.has_plugin(
+                DefaultConnectionHookImpl.__name__
+            ):
+                self.__plugin_registry.load_plugin_with_hooks_impl(
+                    DefaultConnectionHookImpl(), DefaultConnectionHookImpl.__name__
+                )
+            return cast(ConnectionHookSpec, self.__plugin_registry.hook())
+        else:
+            raise PlatformServiceError(
+                ErrorMessage(
+                    "Managed Cursor not initialized properly",
+                    "Cannot connect to database using vdk managed cursor",
+                    "Plugin registry is not initialized. That seems like a bug.",
+                    "Without plugin registry the connection cannot be started",
+                    "Open a vdk github issue "
+                    "and/or revert to previous version of vdk-core.",
+                )
+            )

projects/vdk-core/src/vdk/internal/builtin_plugins/connection/execution_cursor.py

@@ -0,0 +1,67 @@
+# Copyright 2021 VMware, Inc.
+# SPDX-License-Identifier: Apache-2.0
+import logging
+from typing import Any
+
+from vdk.internal.builtin_plugins.connection.decoration_cursor import ManagedOperation
+from vdk.internal.builtin_plugins.connection.pep249.interfaces import PEP249Cursor
+
+
+class ExecuteOperationResult:
+    def __init__(self, native_result: Any):
+        self.__native_result = native_result
+
+    def get_native_result(self) -> Any:
+        """
+        Returns the result of PEP249Cursor.execute invocation using the native cursor.
+        Since this is vendor specific the result type is any - but usually it's number of selected or updated rows.
+        :return: Any
+        """
+        return self.__native_result
+
+
+class ExecutionCursor(PEP249Cursor):
+    """
+    Extends PEP249Cursor to provide:
+        * ability to directly access and execute operations with the native cursor.
+        Generally it should be used only if default implementation does not work (which should be almost never)
+         or more likely - to use some vendor specific features.
+
+    See connection_hook_spec#db_connection_execute_operation for more details and examples how to use it.
+    """
+
+    def __init__(
+        self,
+        native_cursor: PEP249Cursor,
+        managed_operation: ManagedOperation,
+        log=logging.getLogger(__name__),
+    ):
+        super().__init__(native_cursor, log)
+        self.__managed_operation = managed_operation
+        self.__native_cursor = native_cursor
+
+    def get_native_cursor(self) -> PEP249Cursor:
+        """
+        Get the underlying native cursor. Used to actually execute the operation.
+        Generally should not be accessed directly.
+        But it's useful to be exposed in case the native library used provide some extra features.
+        For example Impala (impyla driver) provides a way to get the profile of a query after it's executed
+        which can be printed to aid debugging.
+
+        Check the example usages in docstring of connection_hook_spec#db_connection_execute_operation .
+
+        :return: PEP249Cursor
+            the underlying native cursor being managed.
+        """
+        return self.__native_cursor
+
+    def get_managed_operation(self) -> ManagedOperation:
+        """
+        Retrieve an object that contains information about the query and query parameters used in
+        the database operation. The retrieved Data Transfer Object (DTO) is purposed
+        to curate the query and parameters.
+
+        :return: ManagedOperation
+            Query and parameters DTO
+        """
+        return self.__managed_operation

projects/vdk-core/src/vdk/internal/builtin_plugins/connection/impl/router.py

@@ -5,11 +5,11 @@
 from typing import Dict
 from typing import Union
 
-from vdk.api.plugin.connection_hook_spec import (
-    ConnectionHookSpec,
-)
 from vdk.api.plugin.plugin_input import IManagedConnectionRegistry
 from vdk.internal.builtin_plugins.config.vdk_config import DB_DEFAULT_TYPE
+from vdk.internal.builtin_plugins.connection.connection_hooks import (
+    ConnectionHookSpecFactory,
+)
 from vdk.internal.builtin_plugins.connection.impl.wrapped_connection import (
     WrappedConnection,
 )
@@ -29,9 +29,13 @@ class ManagedConnectionRouter(IManagedConnectionRegistry):
     In both cases dbtype must match the string in which the plugin register itself with.
     """
 
-    def __init__(self, cfg: Configuration, connection_hook_spec: ConnectionHookSpec):
+    def __init__(
+        self,
+        cfg: Configuration,
+        connection_hook_spec_factory: ConnectionHookSpecFactory,
+    ):
         self._cfg: Configuration = cfg
-        self._connection_hook_spec = connection_hook_spec
+        self._connection_hook_spec_factory = connection_hook_spec_factory
         self._log: logging.Logger = logging.getLogger(__name__)
         self._connections: Dict[str, ManagedConnectionBase] = dict()
         self._connection_builders: Dict[
@@ -94,8 +98,8 @@ def __create_connection(self, dbtype):
         conn = self._connection_builders[dbtype]()
         if isinstance(conn, ManagedConnectionBase):
             self._connections[dbtype] = conn
-            if not conn._connection_hook_spec:
-                conn._connection_hook_spec = self._connection_hook_spec
+            if not conn._connection_hook_spec_factory:
+                conn._connection_hook_spec_factory = self._connection_hook_spec_factory
         elif conn is None:
             errors.log_and_throw(
                 to_be_fixed_by=errors.ResolvableBy.CONFIG_ERROR,
@@ -111,6 +115,8 @@ def __create_connection(self, dbtype):
             log = logging.getLogger(conn.__class__.__name__)
             conn.close()  # we will let ManagedConnection to open it when needed.
             self._connections[dbtype] = WrappedConnection(
-                log, self._connection_builders[dbtype], self._connection_hook_spec
+                log,
+                self._connection_builders[dbtype],
+                self._connection_hook_spec_factory,
             )
         return self._connections[dbtype]

projects/vdk-core/src/vdk/internal/builtin_plugins/connection/impl/wrapped_connection.py

@@ -4,8 +4,8 @@
 from typing import Any
 from typing import Callable
 
-from vdk.api.plugin.connection_hook_spec import (
-    ConnectionHookSpec,
+from vdk.internal.builtin_plugins.connection.connection_hooks import (
+    ConnectionHookSpecFactory,
 )
 from vdk.internal.builtin_plugins.connection.managed_connection_base import (
     ManagedConnectionBase,
@@ -22,17 +22,17 @@ def __init__(
         self,
         log: logging.Logger,
         new_connection_builder_function: Callable[[], PEP249Connection],
-        connection_hook_spec: ConnectionHookSpec,
+        connection_hook_spec_factory: ConnectionHookSpecFactory,
     ) -> None:
         """
         :param new_connection_builder_function: method that returns a new (e.g. SAP Hana) connection
             Example
                 def connection() -> ManagedConnectionBase:
                     db = pyhdb.connect(host='hana-prod-d1.northpole.com', port=30015, user='claus', password='hohoho')
                     return db
-        :param connection_hook_spec: connection hook implementations from plugins
+        :param connection_hook_spec_factory: ConnectionHookSpecFactory
         """
-        super().__init__(log, None, connection_hook_spec)
+        super().__init__(log, None, connection_hook_spec_factory)
         self._log = logging.getLogger(__name__)
         self._new_connection_builder_function = new_connection_builder_function
 

projects/vdk-core/src/vdk/internal/builtin_plugins/connection/managed_connection_base.py

@@ -16,8 +16,8 @@
 from tenacity import stop_after_attempt
 from tenacity import wait_exponential
 from vdk.api.job_input import IManagedConnection
-from vdk.api.plugin.connection_hook_spec import (
-    ConnectionHookSpec,
+from vdk.internal.builtin_plugins.connection.connection_hooks import (
+    ConnectionHookSpecFactory,
 )
 from vdk.internal.builtin_plugins.connection.managed_cursor import ManagedCursor
 from vdk.internal.builtin_plugins.connection.pep249.interfaces import PEP249Connection
@@ -40,7 +40,7 @@ def __init__(
         self,
         log: logging.Logger = logger,
         db_con: Optional[PEP249Connection] = None,
-        connection_hook_spec: ConnectionHookSpec = None,
+        connection_hook_spec_factory: ConnectionHookSpecFactory = None,
     ):
         """
         this constructor MUST be called by inheritors
@@ -51,7 +51,7 @@ def __init__(
             self._log = logging.getLogger(__name__)
         self._is_db_con_open: bool = db_con is not None
         self._db_con: Optional[PEP249Connection] = db_con
-        self._connection_hook_spec: ConnectionHookSpec = connection_hook_spec
+        self._connection_hook_spec_factory = connection_hook_spec_factory
 
     def __getattr__(self, attr):
         """
@@ -96,13 +96,14 @@ def method(*args, **kwargs):
     )
     def connect(self) -> PEP249Connection:
         """
-        :return: PEP249 Connection object (unmanaged)
+        :return: PEP249 Connection object (managed)
         """
         if not self._is_db_con_open:
             db_con = self._connect()
             self._log.debug(f"Established {str(db_con)}")
             self._is_db_con_open = True
             self._db_con = db_con
+
         return self
 
     # def get_managed_connection(self) -> PEP249Connection:
@@ -160,7 +161,7 @@ def cursor(self, *args, **kwargs):
             return ManagedCursor(
                 self._db_con.cursor(*args, **kwargs),
                 self._log,
-                self._connection_hook_spec,
+                self._connection_hook_spec_factory,
             )
         return super().cursor()