fdc-shared-kernel 0.0.194__tar.gz → 0.0.196__tar.gz

{fdc_shared_kernel-0.0.194 → fdc_shared_kernel-0.0.196}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: fdc_shared_kernel
-Version: 0.0.194
+Version: 0.0.196
 Summary: Shared library for microservice
 Author-email: Shikhil S <shikhil.s@dbizsolution.com>, Ahammed Akdham N <ahammedakdham.n@dbizsolution.com>
 Classifier: Programming Language :: Python :: 3
@@ -47,6 +47,7 @@ Requires-Dist: snowflake-connector-python==3.16.0
 Requires-Dist: PyMySQL==1.1.1
 Requires-Dist: jaydebeapi==1.2.3
 Requires-Dist: JPype1==1.5.0
+Requires-Dist: psutil==6.0.0
 
 # FDC Shared Kernel
 

{fdc_shared_kernel-0.0.194 → fdc_shared_kernel-0.0.196}/fdc_shared_kernel.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: fdc_shared_kernel
-Version: 0.0.194
+Version: 0.0.196
 Summary: Shared library for microservice
 Author-email: Shikhil S <shikhil.s@dbizsolution.com>, Ahammed Akdham N <ahammedakdham.n@dbizsolution.com>
 Classifier: Programming Language :: Python :: 3
@@ -47,6 +47,7 @@ Requires-Dist: snowflake-connector-python==3.16.0
 Requires-Dist: PyMySQL==1.1.1
 Requires-Dist: jaydebeapi==1.2.3
 Requires-Dist: JPype1==1.5.0
+Requires-Dist: psutil==6.0.0
 
 # FDC Shared Kernel
 

{fdc_shared_kernel-0.0.194 → fdc_shared_kernel-0.0.196}/fdc_shared_kernel.egg-info/SOURCES.txt

@@ -35,6 +35,7 @@ shared_kernel/batch_job/job_storage_managers/s3_job_storage_manager.py
 shared_kernel/config/__init__.py
 shared_kernel/constants/constants.py
 shared_kernel/data_warehouse_handlers/__init__.py
+shared_kernel/data_warehouse_handlers/base_executor.py
 shared_kernel/data_warehouse_handlers/databricks_handler.py
 shared_kernel/data_warehouse_handlers/fabric_sql_handler.py
 shared_kernel/data_warehouse_handlers/factory.py
@@ -183,6 +184,8 @@ shared_kernel/utils/__init__.py
 shared_kernel/utils/data_validators_utils.py
 shared_kernel/utils/date_format_utils.py
 shared_kernel/utils/date_utils.py
+shared_kernel/utils/logging_utils.py
+shared_kernel/utils/request_id.py
 shared_kernel/utils/string_utils.py
 shared_kernel/utils/template_renderer.py
 shared_kernel/utils/thread_debug_utils.py

{fdc_shared_kernel-0.0.194 → fdc_shared_kernel-0.0.196}/fdc_shared_kernel.egg-info/requires.txt

@@ -37,3 +37,4 @@ snowflake-connector-python==3.16.0
 PyMySQL==1.1.1
 jaydebeapi==1.2.3
 JPype1==1.5.0
+psutil==6.0.0

{fdc_shared_kernel-0.0.194 → fdc_shared_kernel-0.0.196}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "fdc_shared_kernel"
-version = "0.0.194"
+version = "0.0.196"
 requires-python = ">=3.7"
 readme = "README_pypi.md"
 description = "Shared library for microservice"

{fdc_shared_kernel-0.0.194 → fdc_shared_kernel-0.0.196}/requirements.txt

@@ -37,3 +37,4 @@ snowflake-connector-python==3.16.0
 PyMySQL==1.1.1
 jaydebeapi==1.2.3
 JPype1==1.5.0
+psutil==6.0.0

fdc_shared_kernel-0.0.196/shared_kernel/data_warehouse_handlers/base_executor.py

@@ -0,0 +1,101 @@
+from abc import ABC
+from functools import wraps
+import logging
+
+from shared_kernel.exceptions.operational_exceptions import QueryCancelledException
+from shared_kernel.interfaces.query_executor import DataWarehouseQueryExecutor
+from shared_kernel.utils.logging_utils import is_request_cancelled
+from shared_kernel.utils.request_id import get_request_id
+
+logger = logging.getLogger(__name__)
+
+def cancel_check(func):
+    """
+    Decorator to wrap query executor execution methods with cancellation checks.
+
+    RATIONALE:
+    When applied to execution methods (e.g. execute_query), it checks if the 
+    request has already been cancelled before sending the query to the warehouse. 
+    It also propagates the distributed request tracing context (request_id) to the executor.
+    """
+    @wraps(func)
+    def wrapper(self, *args, **kwargs):
+        # Capture request ID from the calling thread and store it on self
+        req_id = getattr(self, "request_id", None)
+        if not req_id:
+            req_id = get_request_id()
+            if req_id:
+                self.request_id = req_id
+
+        self.check_cancelled()
+        return func(self, *args, **kwargs)
+    return wrapper
+
+
+class BaseQueryExecutor(DataWarehouseQueryExecutor, ABC):
+    """
+    Base implementation of DataWarehouseQueryExecutor that handles
+    automatic query cancellation checks and raw connection cancellation logic.
+    """
+
+    def check_cancelled(self) -> None:
+        """
+        Checks if the current request has been cancelled, and raises an exception if so.
+        """
+        req_id = getattr(self, "request_id", None)
+        if not req_id:
+            req_id = get_request_id()
+
+        if req_id and is_request_cancelled(req_id):
+            logger.warning(f"[REQ_ID: {req_id}] Aborting query execution due to client disconnect.")
+            raise QueryCancelledException("Query execution aborted due to client disconnect.")
+
+    @property
+    def active_connections(self) -> set:
+        if not hasattr(self, '_active_connections'):
+            self._active_connections = set()
+        return self._active_connections
+
+    def _cancel_connection(self, conn) -> None:
+        """
+        Generic helper to cancel an individual active connection.
+        Handles SQLAlchemy connection wrappers, DBAPI connections, and fallback logic.
+
+        RATIONALE:
+        1. We attempt to invoke driver-native cancel() (e.g. raw_conn.cancel()) which cancels the 
+           active query safely on the database server.
+        2. If native cancel is not supported by the driver, we log a warning and DO NOTHING (pass). 
+           This allows the query to run to completion or fail gracefully on timeout, ensuring the 
+           connection remains in a healthy state and is safely returned back to the pool.
+        """
+        cancelled = False
+        try:
+            raw_conn = None
+            if hasattr(conn, 'connection') and conn.connection:
+                # conn.connection is typically a _ConnectionFairy in SQLAlchemy
+                fairy = conn.connection
+                if hasattr(fairy, 'dbapi_connection') and fairy.dbapi_connection:
+                    raw_conn = fairy.dbapi_connection
+                elif hasattr(fairy, 'connection') and fairy.connection:
+                    raw_conn = fairy.connection
+                else:
+                    raw_conn = fairy
+            elif hasattr(conn, 'dbapi_connection') and conn.dbapi_connection:
+                raw_conn = conn.dbapi_connection
+            else:
+                raw_conn = conn
+
+            if raw_conn and hasattr(raw_conn, 'cancel'):
+                logger.info("Calling driver-level cancel() on raw connection.")
+                raw_conn.cancel()
+                cancelled = True
+            elif hasattr(conn, 'cancel'):
+                logger.info("Calling cancel() on connection.")
+                conn.cancel()
+                cancelled = True
+        except Exception as cancel_err:
+            logger.error(f"Error calling driver-level cancel: {cancel_err}")
+
+        if not cancelled:
+            logger.warning("Driver does not support cancel(). Doing nothing.")
+            pass

{fdc_shared_kernel-0.0.194 → fdc_shared_kernel-0.0.196}/shared_kernel/data_warehouse_handlers/connections/databricks_warehouse_connection.py

@@ -7,7 +7,7 @@ from databricks.sql.exc import InterfaceError, OperationalError, ProgrammingErro
 
 from shared_kernel.dataclasses.warehouse_configs import DatabricksConfig
 from shared_kernel.exceptions.http_exceptions import BadRequest
-from shared_kernel.interfaces.warehouse_connection import DataWarehouseConnection
+from shared_kernel.interfaces.warehouse_connection import DataWarehouseConnection, cancel_check_connection
 from shared_kernel.logger import Logger
 from shared_kernel.config import Config
 
@@ -88,6 +88,7 @@ class DatabricksWarehouseConnection(DataWarehouseConnection):
                 self.close_connection()
 
     @contextmanager
+    @cancel_check_connection
     def get_connection(self, retries: int = 3, delay: int = 5, query_timeout: int = 30):
         if not (self._conn and self.is_valid_connection()):
             self._create_connection_with_retry(retries, delay)

{fdc_shared_kernel-0.0.194 → fdc_shared_kernel-0.0.196}/shared_kernel/data_warehouse_handlers/connections/fabric_sql_warehouse_connection.py

@@ -14,7 +14,7 @@ from sqlalchemy.exc import InterfaceError, OperationalError, ProgrammingError
 
 from shared_kernel.dataclasses.warehouse_configs import FabricSQLConfig
 from shared_kernel.exceptions.http_exceptions import BadRequest
-from shared_kernel.interfaces.warehouse_connection import DataWarehouseConnection
+from shared_kernel.interfaces.warehouse_connection import DataWarehouseConnection, cancel_check_connection
 from shared_kernel.data_warehouse_handlers.utils import extract_sql_server_error_message
 from shared_kernel.logger import Logger
 from shared_kernel.config import Config
@@ -85,12 +85,13 @@ class FabricSQLWarehouseConnection(DataWarehouseConnection):
                     from sqlalchemy import event
                     
                     connection_string = self._get_fabric_connection_params(cfg)
+                    pool_size = int(Config.get("FABRIC_SQL_POOL_SIZE", 15))
 
                     engine = create_engine(
                         connection_string,
-                        pool_size=50,
+                        pool_size=pool_size,
                         max_overflow=5,
-                        pool_timeout=60,
+                        pool_timeout=300,
                         pool_recycle=1800,
                         pool_pre_ping=True,
                     )
@@ -188,6 +189,7 @@ class FabricSQLWarehouseConnection(DataWarehouseConnection):
         return self._engines[key]
 
     @contextmanager
+    @cancel_check_connection
     def get_connection(self, retries: int = 3, delay: int = 5, query_timeout: int = 30):
         """Context manager for database connections with retry mechanism"""
         engine = self._get_engine()

{fdc_shared_kernel-0.0.194 → fdc_shared_kernel-0.0.196}/shared_kernel/data_warehouse_handlers/connections/jdbc_warehouse_connection.py

@@ -9,7 +9,7 @@ from typing import Dict, Any, Optional
 
 from shared_kernel.dataclasses.warehouse_configs import JDBCConfig
 from shared_kernel.exceptions.http_exceptions import BadRequest
-from shared_kernel.interfaces.warehouse_connection import DataWarehouseConnection
+from shared_kernel.interfaces.warehouse_connection import DataWarehouseConnection, cancel_check_connection
 from shared_kernel.logger import Logger
 from shared_kernel.config import Config
 from shared_kernel.data_warehouse_handlers.jdbc_driver_registry import get_jdbc_driver_config
@@ -112,6 +112,7 @@ class JDBCWarehouseConnection(DataWarehouseConnection):
                 # We don't raise yet, as the class might somehow be in the global classpath
 
     @contextmanager
+    @cancel_check_connection
     def get_connection(self, retries: int = 3, delay: int = 5, query_timeout: int = 30):
         """Context manager for JDBC database connections with retry mechanism."""
         last_exception = None

{fdc_shared_kernel-0.0.194 → fdc_shared_kernel-0.0.196}/shared_kernel/data_warehouse_handlers/connections/mssql_warehouse_connection.py

@@ -12,7 +12,7 @@ from sqlalchemy.exc import InterfaceError, OperationalError, ProgrammingError
 
 from shared_kernel.dataclasses.warehouse_configs import MSSQLConfig
 from shared_kernel.exceptions.http_exceptions import BadRequest
-from shared_kernel.interfaces.warehouse_connection import DataWarehouseConnection
+from shared_kernel.interfaces.warehouse_connection import DataWarehouseConnection, cancel_check_connection
 from shared_kernel.constants.constants import get_constant
 from shared_kernel.data_warehouse_handlers.utils import extract_sql_server_error_message
 from shared_kernel.logger import Logger
@@ -62,10 +62,12 @@ class MSSQLWarehouseConnection(DataWarehouseConnection):
                     )
                     connect_args = {}
 
+                    pool_size = int(Config.get("MSSQL_POOL_SIZE", 10))
+
                     self._engines[key] = create_engine(
                         connection_string,
                         connect_args=connect_args,
-                        pool_size=10,
+                        pool_size=pool_size,
                         max_overflow=5,
                         pool_timeout=30,
                         pool_recycle=1800,
@@ -78,6 +80,7 @@ class MSSQLWarehouseConnection(DataWarehouseConnection):
         return self._engines[key]
 
     @contextmanager
+    @cancel_check_connection
     def get_connection(self, retries: int = 3, delay: int = 5, query_timeout: int = 30):
         """Context manager for database connections with retry mechanism"""
         engine = self._get_engine()

{fdc_shared_kernel-0.0.194 → fdc_shared_kernel-0.0.196}/shared_kernel/data_warehouse_handlers/connections/mysql_warehouse_connection.py

@@ -14,7 +14,7 @@ from sqlalchemy.exc import InterfaceError, OperationalError, ProgrammingError
 from shared_kernel.constants.constants import get_constant
 from shared_kernel.dataclasses.warehouse_configs import MySQLConfig
 from shared_kernel.exceptions.http_exceptions import BadRequest
-from shared_kernel.interfaces.warehouse_connection import DataWarehouseConnection
+from shared_kernel.interfaces.warehouse_connection import DataWarehouseConnection, cancel_check_connection
 from shared_kernel.logger import Logger
 from shared_kernel.config import Config
 
@@ -68,10 +68,11 @@ class MySQLWarehouseConnection(DataWarehouseConnection):
                     connection_string = get_constant("mysql", "CONNECTION_STRING").format(
                         cfg.username, password, cfg.host, int(cfg.port), cfg.database
                     )
+                    pool_size = int(Config.get("MYSQL_POOL_SIZE", 10))
                     self._engines[key] = create_engine(
                         connection_string,
                         connect_args=self._CONNECT_ARGS,
-                        pool_size=10,
+                        pool_size=pool_size,
                         max_overflow=5,
                         pool_timeout=30,
                         pool_recycle=1800,
@@ -82,6 +83,7 @@ class MySQLWarehouseConnection(DataWarehouseConnection):
         return self._engines[key]
 
     @contextmanager
+    @cancel_check_connection
     def get_connection(self, retries: int = 3, delay: int = 5, query_timeout: int = 30):
         engine = self._get_engine()
         last_exception = None

{fdc_shared_kernel-0.0.194 → fdc_shared_kernel-0.0.196}/shared_kernel/data_warehouse_handlers/connections/postgresql_warehouse_connection.py

@@ -12,7 +12,7 @@ from sqlalchemy.exc import InterfaceError, OperationalError, ProgrammingError, S
 
 from shared_kernel.dataclasses.warehouse_configs import PostgreSQLConfig
 from shared_kernel.exceptions.http_exceptions import BadRequest
-from shared_kernel.interfaces.warehouse_connection import DataWarehouseConnection
+from shared_kernel.interfaces.warehouse_connection import DataWarehouseConnection, cancel_check_connection
 from shared_kernel.logger import Logger
 from shared_kernel.config import Config
 
@@ -61,9 +61,10 @@ class PostgreSQLWarehouseConnection(DataWarehouseConnection):
                         f"postgresql+psycopg2://{cfg.username}:{password}"
                         f"@{cfg.host}:{cfg.port}/{cfg.database}"
                     )
+                    pool_size = int(Config.get("POSTGRES_POOL_SIZE", 10))
                     self._engines[key] = create_engine(
                         url,
-                        pool_size=10,
+                        pool_size=pool_size,
                         max_overflow=5,
                         pool_timeout=30,
                         pool_recycle=1800,
@@ -75,6 +76,7 @@ class PostgreSQLWarehouseConnection(DataWarehouseConnection):
         return self._engines[key]
 
     @contextmanager
+    @cancel_check_connection
     def get_connection(self, retries: int = 3, delay: int = 5, query_timeout: int = 30):
         """Context manager for database connections with retry mechanism"""
         engine = self._get_engine()

{fdc_shared_kernel-0.0.194 → fdc_shared_kernel-0.0.196}/shared_kernel/data_warehouse_handlers/connections/redshift_warehouse_connection.py

@@ -12,7 +12,7 @@ from sqlalchemy.exc import InterfaceError, OperationalError, ProgrammingError
 
 from shared_kernel.dataclasses.warehouse_configs import RedshiftConfig
 from shared_kernel.exceptions.http_exceptions import BadRequest
-from shared_kernel.interfaces.warehouse_connection import DataWarehouseConnection
+from shared_kernel.interfaces.warehouse_connection import DataWarehouseConnection, cancel_check_connection
 from shared_kernel.logger import Logger
 from shared_kernel.config import Config
 
@@ -194,9 +194,10 @@ class RedshiftWarehouseConnection(DataWarehouseConnection):
                         f"redshift+redshift_connector://{cfg.username}:{cfg.password}"
                         f"@{cfg.host}:{cfg.port}/{cfg.database}"
                     )
+                    pool_size = int(Config.get("REDSHIFT_POOL_SIZE", 10))
                     self._engines[key] = create_engine(
                         url,
-                        pool_size=10,
+                        pool_size=pool_size,
                         max_overflow=5,
                         pool_timeout=30,
                         pool_recycle=1800,
@@ -222,6 +223,7 @@ class RedshiftWarehouseConnection(DataWarehouseConnection):
         return False
 
     @contextmanager
+    @cancel_check_connection
     def get_connection(self, retries: int = 3, delay: int = 5, query_timeout: int = 30):
         """Context manager for database connections with retry mechanism"""
         engine = self._get_engine()

{fdc_shared_kernel-0.0.194 → fdc_shared_kernel-0.0.196}/shared_kernel/data_warehouse_handlers/connections/snowflake_warehouse_connection.py

@@ -19,7 +19,7 @@ from snowflake.connector.errors import (
 from shared_kernel.data_warehouse_handlers.utils import extract_snowflake_error_message
 from shared_kernel.dataclasses.warehouse_configs import SnowflakeConfig
 from shared_kernel.exceptions.http_exceptions import BadRequest
-from shared_kernel.interfaces.warehouse_connection import DataWarehouseConnection
+from shared_kernel.interfaces.warehouse_connection import DataWarehouseConnection, cancel_check_connection
 from shared_kernel.logger import Logger
 from shared_kernel.config import Config
 
@@ -111,6 +111,7 @@ class SnowflakeWarehouseConnection(DataWarehouseConnection):
                 self.close_connection()
 
     @contextmanager
+    @cancel_check_connection
     def get_connection(self, retries: int = 3, delay: int = 5, query_timeout: int = 30):
         if not (self._conn and self.is_valid_connection()):
             self._create_connection_with_retry(retries, delay)

{fdc_shared_kernel-0.0.194 → fdc_shared_kernel-0.0.196}/shared_kernel/data_warehouse_handlers/query_executors/databricks_query_executor.py

@@ -6,7 +6,7 @@ from typing import Any, Dict, List, Optional, Tuple
 from databricks.sql.exc import InterfaceError, OperationalError, ProgrammingError
 
 from shared_kernel.config import Config
-from shared_kernel.interfaces.query_executor import DataWarehouseQueryExecutor
+from shared_kernel.data_warehouse_handlers.base_executor import BaseQueryExecutor, cancel_check
 from shared_kernel.interfaces.warehouse_connection import DataWarehouseConnection
 from shared_kernel.logger import Logger
 
@@ -15,11 +15,12 @@ logger = Logger(config.get("APP_NAME"))
 MAX_QUERY_CONCURRENCY = 5
 
 
-class DatabricksQueryExecutor(DataWarehouseQueryExecutor):
+class DatabricksQueryExecutor(BaseQueryExecutor):
 
     def __init__(self, warehouse_connection: DataWarehouseConnection) -> None:
         self.warehouse_connection = warehouse_connection
 
+    @cancel_check
     def execute_query(
         self,
         query: str,
@@ -29,6 +30,7 @@ class DatabricksQueryExecutor(DataWarehouseQueryExecutor):
         Execute a single SQL query and return all rows as a list of dicts.
         """
         with self.warehouse_connection.get_connection() as cursor:
+            self.active_connections.add(cursor)
             try:
                 logger.info(f"Executing query: {query}")
                 start = time.perf_counter()
@@ -49,6 +51,8 @@ class DatabricksQueryExecutor(DataWarehouseQueryExecutor):
                 raise Exception(f"Query execution failed: {str(e)}") from e
             except Exception as e:
                 raise Exception(f"Query execution failed: {str(e)}") from e
+            finally:
+                self.active_connections.discard(cursor)
 
     def execute_queries(
         self,
@@ -84,12 +88,14 @@ class DatabricksQueryExecutor(DataWarehouseQueryExecutor):
 
         return results
     
+    @cancel_check
     def execute_query_with_column_types(
         self,
         query: str,
         params: Optional[tuple] = None,
     ) -> Tuple[List[Dict[str, Any]], List[Tuple]]:
         with self.warehouse_connection.get_connection() as cursor:
+            self.active_connections.add(cursor)
             try:
                 logger.info(f"Executing query with column types: {query}")
                 start = time.perf_counter()
@@ -115,4 +121,14 @@ class DatabricksQueryExecutor(DataWarehouseQueryExecutor):
                 raise Exception(f"Query execution failed: {str(e)}") from e
             except Exception as e:
                 raise Exception(f"Query execution failed: {str(e)}") from e
+            finally:
+                self.active_connections.discard(cursor)
+
+    def cancel(self) -> None:
+        """
+        Cancels all active Databricks query cursors.
+        """
+        logger.warning(f"Cancelling {len(self.active_connections)} active Databricks query cursor(s).")
+        for cursor in list(self.active_connections):
+            self._cancel_connection(cursor)
     

{fdc_shared_kernel-0.0.194 → fdc_shared_kernel-0.0.196}/shared_kernel/data_warehouse_handlers/query_executors/fabric_sql_query_executor.py

@@ -8,7 +8,7 @@ from sqlalchemy.exc import InterfaceError, OperationalError, ProgrammingError
 
 from shared_kernel.config import Config
 from shared_kernel.data_warehouse_handlers.utils import extract_error_message
-from shared_kernel.interfaces.query_executor import DataWarehouseQueryExecutor
+from shared_kernel.data_warehouse_handlers.base_executor import BaseQueryExecutor, cancel_check
 from shared_kernel.interfaces.warehouse_connection import DataWarehouseConnection
 from shared_kernel.logger import Logger
 
@@ -17,7 +17,7 @@ logger = Logger(config.get("APP_NAME"))
 MAX_QUERY_CONCURRENCY = 5
 
 
-class FabricSQLQueryExecutor(DataWarehouseQueryExecutor):
+class FabricSQLQueryExecutor(BaseQueryExecutor):
     """
     Query executor for Microsoft Fabric SQL Server
     """
@@ -25,6 +25,7 @@ class FabricSQLQueryExecutor(DataWarehouseQueryExecutor):
     def __init__(self, warehouse_connection: DataWarehouseConnection) -> None:
         self.warehouse_connection = warehouse_connection
 
+    @cancel_check
     def execute_query(
         self,
         query: str,
@@ -34,15 +35,12 @@ class FabricSQLQueryExecutor(DataWarehouseQueryExecutor):
         Execute a single SQL query and return all rows as a list of dicts.
         """
         with self.warehouse_connection.get_connection() as conn:
+            self.active_connections.add(conn)
             try:
                 logger.info(f"Executing query: {query}")
                 start = time.perf_counter()
 
-                result = (
-                    conn.execute(text(query), params)
-                    if params
-                    else conn.execute(text(query))
-                )
+                result = conn.execute(text(query), params) if params else conn.execute(text(query))
 
                 elapsed = time.perf_counter() - start
                 logger.info(f"Query executed in {elapsed:.6f}s")
@@ -64,6 +62,8 @@ class FabricSQLQueryExecutor(DataWarehouseQueryExecutor):
             except Exception as e:
                 user_msg = extract_error_message(str(e))
                 raise Exception(f"Query execution failed: {user_msg}") from e
+            finally:
+                self.active_connections.discard(conn)
 
     def execute_queries(
         self,
@@ -99,6 +99,7 @@ class FabricSQLQueryExecutor(DataWarehouseQueryExecutor):
 
         return results
 
+    @cancel_check
     def execute_query_with_column_types(
         self,
         query: str,
@@ -108,15 +109,12 @@ class FabricSQLQueryExecutor(DataWarehouseQueryExecutor):
         Execute a query and return both rows and raw DBAPI cursor description.
         """
         with self.warehouse_connection.get_connection() as conn:
+            self.active_connections.add(conn)
             try:
                 logger.info(f"Executing query with column types: {query}")
                 start = time.perf_counter()
 
-                result = (
-                    conn.execute(text(query), params)
-                    if params
-                    else conn.execute(text(query))
-                )
+                result = conn.execute(text(query), params) if params else conn.execute(text(query))
 
                 elapsed = time.perf_counter() - start
                 logger.info(f"Query executed in {elapsed:.6f}s")
@@ -140,3 +138,13 @@ class FabricSQLQueryExecutor(DataWarehouseQueryExecutor):
             except Exception as e:
                 user_msg = extract_error_message(str(e))
                 raise Exception(f"Query execution failed: {user_msg}") from e
+            finally:
+                self.active_connections.discard(conn)
+
+    def cancel(self) -> None:
+        """
+        Cancels all active Fabric SQL query connections.
+        """
+        logger.warning(f"Cancelling {len(self.active_connections)} active Fabric SQL query connection(s).")
+        for conn in list(self.active_connections):
+            self._cancel_connection(conn)

{fdc_shared_kernel-0.0.194 → fdc_shared_kernel-0.0.196}/shared_kernel/data_warehouse_handlers/query_executors/jdbc_query_executor.py

@@ -3,7 +3,7 @@ from concurrent.futures import ThreadPoolExecutor, as_completed
 from typing import Any, Dict, List, Optional, Tuple
 
 from shared_kernel.config import Config
-from shared_kernel.interfaces.query_executor import DataWarehouseQueryExecutor
+from shared_kernel.data_warehouse_handlers.base_executor import BaseQueryExecutor, cancel_check
 from shared_kernel.interfaces.warehouse_connection import DataWarehouseConnection
 from shared_kernel.logger import Logger
 
@@ -11,7 +11,7 @@ config = Config()
 logger = Logger(config.get("APP_NAME"))
 MAX_QUERY_CONCURRENCY = 5
 
-class JDBCQueryExecutor(DataWarehouseQueryExecutor):
+class JDBCQueryExecutor(BaseQueryExecutor):
     """
     Query executor for JDBC connections.
     """
@@ -19,6 +19,7 @@ class JDBCQueryExecutor(DataWarehouseQueryExecutor):
     def __init__(self, warehouse_connection: DataWarehouseConnection) -> None:
         self.warehouse_connection = warehouse_connection
 
+    @cancel_check
     def execute_query(
         self,
         query: str,
@@ -28,6 +29,7 @@ class JDBCQueryExecutor(DataWarehouseQueryExecutor):
         Execute a single SQL query and return all rows as a list of dicts.
         """
         with self.warehouse_connection.get_connection() as conn:
+            self.active_connections.add(conn)
             cursor = conn.cursor()
             try:
                 logger.info(f"Executing JDBC query: {query}")
@@ -56,6 +58,7 @@ class JDBCQueryExecutor(DataWarehouseQueryExecutor):
                 raise Exception(f"JDBC Query execution failed: {str(e)}") from e
             finally:
                 cursor.close()
+                self.active_connections.discard(conn)
 
     def execute_queries(
         self,
@@ -87,6 +90,7 @@ class JDBCQueryExecutor(DataWarehouseQueryExecutor):
 
         return results
 
+    @cancel_check
     def execute_query_with_column_types(
         self,
         query: str,
@@ -96,6 +100,7 @@ class JDBCQueryExecutor(DataWarehouseQueryExecutor):
         Execute a query and return both rows and raw cursor description.
         """
         with self.warehouse_connection.get_connection() as conn:
+            self.active_connections.add(conn)
             cursor = conn.cursor()
             try:
                 logger.info(f"Executing JDBC query with column types: {query}")
@@ -125,3 +130,16 @@ class JDBCQueryExecutor(DataWarehouseQueryExecutor):
                 raise Exception(f"JDBC Query execution failed: {str(e)}") from e
             finally:
                 cursor.close()
+                self.active_connections.discard(conn)
+
+    def cancel(self) -> None:
+        """
+        Cancels all active JDBC query connections by closing them directly.
+        """
+        logger.warning(f"Cancelling {len(self.active_connections)} active JDBC connection(s).")
+        for conn in list(self.active_connections):
+            try:
+                logger.info("Closing active JDBC connection to abort running query.")
+                conn.close()
+            except Exception as e:
+                logger.error(f"Error closing JDBC connection during cancel: {e}")