NucelarMonitor/Python/Drivers/SQLServerDriver.py

121 lines
4.6 KiB
Python

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
from typing import Self, Any, Optional, Sequence
from re import Match as REMatch
from Assets.pyodbc import Connection as Connection, connect as connect, Cursor as Cursor
from Interfaces.Application.NucelarMonitorInterface import NucelarMonitorInterface
from Abstracts.DatabaseAbstract import DatabaseAbstract
from Models.QueryResponseModel import QueryResponseModel
from Utils.Utils import Utils
from Utils.Patterns import RE
class SQLServerDriver(DatabaseAbstract):
def __init__(self:Self,
nucelar_monitor:NucelarMonitorInterface,
inputs:Optional[dict[str, Any|None]|Sequence[Any|None]] = None
) -> None:
super().__init__(nucelar_monitor, inputs)
self.__connection:Connection|None = None
self.__string_connection:str = Utils.string_variables(self.nucelar_monitor.settings.get((
"odbc_string_connection", "sql_string_connection", "string_connection"
), inputs, "DRIVER={driver};SERVER={host},{port};UID={user};PWD={password};DATABASE={database}"), {
"driver" : self.nucelar_monitor.settings.get("sql_driver", inputs, "{ODBC Driver 17 for SQL Server}"),
"host" : self.nucelar_monitor.settings.get(("sql_host", "odbc_host"), inputs, "localhost"),
"port" : self.nucelar_monitor.settings.get(("sql_port", "odbc_port"), inputs, 1433),
"user" : self.nucelar_monitor.settings.get(("sql_user", "odbc_user"), inputs, "sa"),
"password" : self.nucelar_monitor.settings.get(("sql_password", "odbc_password"), inputs, "password"),
"database" : self.nucelar_monitor.settings.get(("sql_database", "odbc_database"), inputs, "NucelarMonitor")
})
def close(self:Self) -> bool:
if self.__connection is not None:
try:
self.__connection.close()
return True
except Exception as exception:
self.nucelar_monitor.exception(exception, "sql_server_close_exception", {
"string_connection" : self.__string_connection
})
return False
return True
def __autoclose(self:Self) -> None:
pass
def format_query(self:Self,
query:str,
parameters:Optional[dict[str, Any|None]|Sequence[Any|None]] = None
) -> tuple[str, list[str]]:
variables:list[str] = []
def callback(matches:REMatch) -> str:
key:str = matches.group(1)
if key:
return (
"'" + str(parameters[key]).replace("'","''") + "'" if isinstance(parameters[key], str) else
str(parameters[key]) if key in parameters else matches.group(0))
key = matches.group(2)
variables.append(key)
return matches.group(0)
query = RE.ODBC_STRING_VARIABLE.sub(callback, query)
return (
"".join("declare @" + variable + " varchar(max)\n" for variable in variables) +
query +
"\nselect " + ", ".join("@" + variable for variable in variables)
) if len(variables) else query, variables
def query(self:Self,
query:str,
parameters:Optional[dict[str, Any|None]|Sequence[Any|None]] = None
) -> QueryResponseModel:
response:QueryResponseModel = QueryResponseModel()
cursor:Cursor|None = None
variables:list[str] = []
try:
if self.__connection is None:
self.__connection = connect(
self.__string_connection,
autocommit = True
)
cursor = self.__connection.cursor()
query, variables = self.format_query(query, parameters)
cursor.execute(query)
while True:
if cursor.description is not None:
response.columns.append([column[0] for column in cursor.description])
response.tables.append([tuple(row) for row in cursor.fetchall()])
if not cursor.nextset():
break
except Exception as exception:
self.nucelar_monitor.exception(exception, "connection_exception", {
"string_connection" : self.__string_connection
})
finally:
if cursor is not None:
cursor.close()
if len(variables) and len(response.tables):
for i, variable in enumerate(variables):
response.variables[variable] = response.tables[-1][0][i]
response.tables = response.tables[:-1]
response.columns = response.columns[:-1]
return response