import logging
from argparse import _SubParsersAction
from typing import Optional, Type, Union
from lxml.etree import ElementBase
from q2_sdk.core.configuration import settings
from q2_sdk.core.default_settings import RUNNING_ENTRYPOINT
from q2_sdk.core.dynamic_imports import (
api_ExecuteStoredProcedure,
ob_ExecuteStoredProcedure,
)
from q2_sdk.core.exceptions import DatabaseDataError, HqResponseError
from q2_sdk.core.opentelemetry.span import Q2Span, Q2SpanAttributes
from q2_sdk.hq.models.hq_credentials import HqCredentials
from q2_sdk.hq.models.hq_response import HqResponse
from q2_sdk.hq.table_row import TableRow
from q2_sdk.tools.sentinel import Sentinel
from q2_sdk.tools.utils import serialize_for_cli
DEFAULT_LOGGER = logging.getLogger()
DEFAULT = Sentinel("DEFAULT_DB_PARAM")
MISSING_AUDIT_ACTION_TABLE = "Q2_AuditAction"
UNREGISTERED_AUDIT_ACTIONS: set[str] = set()
[docs]
class DbObject:
GET_BY_NAME_KEY = ""
NAME = ""
REPRESENTATION_ROW_CLASS: Optional[Type[TableRow]] = None
AUDIT_EXEMPT_STORED_PROCS: frozenset[str] = frozenset()
def __init__(
self,
logger,
hq_credentials: Optional[HqCredentials] = None,
ret_table_obj: Optional[bool] = None,
):
"""
Programmatic access to the Q2 database. Not as flexible as a true ORM, but takes the
guesswork out of database schemas and ensures safety in the transactions.
:param logger: Reference to calling request's logger (self.logger in your extension)
:param hq_credentials: HQ Connectivity Information (Defaults to settings file)
:param ret_table_obj: Flag to return list of LXML elements if ``False`` or
TableRow objects from DB calls if ``True`` (Defaults to settings file)
"""
if ret_table_obj is None:
ret_table_obj = settings.RETURN_TABLE_OBJECTS_FROM_DB
if not logger:
logger = DEFAULT_LOGGER
self.logger = logger
self.hq_response: Optional[HqResponse] = None
self.ret_table_obj = ret_table_obj
self._hq_credentials = hq_credentials if hq_credentials else None
[docs]
def add_arguments(self, parser: _SubParsersAction):
"""
Hook for subclassed DbObjects to add custom arguments.
"""
@property
def hq_credentials(self) -> HqCredentials:
if not self._hq_credentials:
self._hq_credentials = settings.HQ_CREDENTIALS
return self._hq_credentials
[docs]
@Q2Span.instrument(skip=["stored_proc_short_name"])
async def call_hq(
self,
stored_proc_short_name: str,
sql_parameters: Optional[
Union[
api_ExecuteStoredProcedure.SqlParameters,
ob_ExecuteStoredProcedure.SqlParameters,
]
] = None,
specific_table: str = "Table",
representation_class_override=None,
use_json=True,
force_q2_api=False,
**kwargs,
):
"""
:param stored_proc_short_name: Registered in the Q2_ApiStoredProc table
:param sql_parameters: Params to the stored proc
:param specific_table: If shape returned is different than standard (rare)
:param representation_class_override: Used in cases where a DbObject has more than one return class defined
:param use_json: If False, will use the pure soap interface
:param force_q2_api: If True, will not use wedge_online_banking, even if there is an active session
"""
if sql_parameters:
for param in sql_parameters.sql_param:
if isinstance(param.value, ElementBase):
param.value = param.value.pyval
Q2Span.set_attribute(
Q2SpanAttributes.STORED_PROCEDURE_NAME, stored_proc_short_name
)
representation_row_class = (
representation_class_override or self.REPRESENTATION_ROW_CLASS
)
try:
self.hq_response = await self._call_execute_stored_procedure(
stored_proc_short_name,
sql_parameters,
use_json=use_json,
force_q2_api=force_q2_api,
**kwargs,
)
except Exception as error:
await self._audit_stored_proc_execution(stored_proc_short_name, str(error))
raise
if self.hq_response.success is False:
await self._audit_stored_proc_execution(
stored_proc_short_name, self.hq_response.error_message
)
raise HqResponseError(
f'HQ Request returned with error message: "{self.hq_response.error_message}"'
)
await self._audit_stored_proc_execution(stored_proc_short_name)
if self.ret_table_obj:
return self.hq_response.parse_sproc_return(
representation_row_class=representation_row_class,
specific_table=specific_table,
)
else:
return self.hq_response.parse_sproc_return(specific_table=specific_table)
async def _audit_stored_proc_execution(
self, stored_proc_short_name: str, exception_message: Optional[str] = None
) -> None:
"""
HQ audits every execution under its generic ExecuteStoredProc action, which says nothing about which proc
actually ran, so this writes a record keyed to the proc's own name. Procs with no dedicated audit action are
left to HQ's record: the first execution asks, and HQ's rejection is remembered for the rest of the process,
so the behavior follows whatever is actually registered in this environment.
Auditing must never break the database call it describes, so failures are logged rather than raised.
:param stored_proc_short_name: Registered in the Q2_ApiStoredProc table
:param exception_message: Reason the execution failed, if it did
"""
if any((
not settings.AUDIT_STORED_PROC_EXECUTIONS,
stored_proc_short_name in self.AUDIT_EXEMPT_STORED_PROCS,
stored_proc_short_name in UNREGISTERED_AUDIT_ACTIONS,
)):
return
from q2_sdk.hq.db.audit_record import AuditRecord
try:
audit = AuditRecord(self.logger, self.hq_credentials)
response = await audit.create(
f"SDK stored procedure execution: {stored_proc_short_name}",
session_id="-1",
audit_action_short_name=stored_proc_short_name,
exception_message=exception_message,
)
if response.success is False:
if MISSING_AUDIT_ACTION_TABLE in (response.error_message or ""):
UNREGISTERED_AUDIT_ACTIONS.add(stored_proc_short_name)
self.logger.debug(
"%s has no dedicated audit action, leaving it to HQ's ExecuteStoredProc record",
stored_proc_short_name,
)
else:
self.logger.warning(
"HQ rejected the audit record for stored procedure %s: %s",
stored_proc_short_name,
response.error_message,
)
except Exception:
self.logger.warning(
"Could not write audit record for stored procedure %s",
stored_proc_short_name,
exc_info=True,
)
async def _call_execute_stored_procedure(
self,
stored_proc_short_name: str,
sql_parameters: Optional[
Union[
api_ExecuteStoredProcedure.SqlParameters,
ob_ExecuteStoredProcedure.SqlParameters,
]
] = None,
use_json=True,
force_q2_api=False,
**kwargs,
) -> HqResponse:
if self.hq_credentials.auth_token and not force_q2_api:
execute_module = ob_ExecuteStoredProcedure
else:
execute_module = api_ExecuteStoredProcedure
if RUNNING_ENTRYPOINT == "run":
self.logger.info(f"Executing Stored Procedure: {stored_proc_short_name}")
result = await execute_module.execute(
execute_module.ParamsObj(
logger=self.logger,
stored_proc_short_name=stored_proc_short_name,
sql_parameters=sql_parameters,
hq_credentials=self.hq_credentials,
),
use_json=use_json,
**kwargs,
)
return result
[docs]
@staticmethod
def serialize_for_cli(
rows: list,
fields_to_display: Optional[list[str]] = None,
fields_to_truncate: Optional[list[str]] = None,
):
"""
Tab delimits response for printing to a terminal
:param rows: XML elements from HQ response
:param fields_to_display: Optional. Displays all fields otherwise
:param fields_to_truncate: Optional. Limits display of these fields to 15 characters
:return: Tab delimited database rows
"""
return serialize_for_cli(rows, fields_to_display, fields_to_truncate)
[docs]
async def get_by_name(self, name, get_by_name_key=None, get_func=None, **kwargs):
if not get_by_name_key:
get_by_name_key = self.GET_BY_NAME_KEY
if not get_func:
get_func = self.get
if not (get_by_name_key and self.NAME):
raise NotImplementedError
full_response = await get_func(**kwargs)
filtered_response = [
x for x in full_response if x.findtext(get_by_name_key) == name
]
if not filtered_response:
raise DatabaseDataError(f'No {self.NAME} with name "{name}"')
else:
if len(filtered_response) > 1:
raise DatabaseDataError(f'More than one {self.NAME} with name "{name}"')
product = filtered_response[0]
return product