khushalps
3c4cac1019
git-subtree-dir: utils_v2 git-subtree-split: af73d53e43729f79a736602775d61ef5b1b0d9cf
437 lines
16 KiB
Python
437 lines
16 KiB
Python
"""
|
|
|
|
AUTHOR:
|
|
|
|
Khushal P Soonderji
|
|
|
|
DATE:
|
|
|
|
Friday, 30th Aug., 2024
|
|
|
|
OBJECTIVE:
|
|
|
|
To be able to access SQL-based databases from python in a simple way.
|
|
|
|
REFERENCES:
|
|
|
|
N/A
|
|
|
|
DOWNLOADS:
|
|
|
|
N/A
|
|
|
|
"""
|
|
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** IMPORT ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
# To make sibling directories accessible for imports:
|
|
import sys
|
|
sys.path.append(".")
|
|
sys.path.append("..")
|
|
|
|
# MySQL Database:
|
|
import aiomysql
|
|
import decimal
|
|
|
|
# For data-crunching:
|
|
import pandas as pd
|
|
|
|
# For time-keeping:
|
|
import time
|
|
|
|
# OS-level operations:
|
|
import os
|
|
|
|
# My utils:
|
|
from utils_v2.string import json
|
|
|
|
# For async activities:
|
|
import asyncio
|
|
|
|
# For debugging:
|
|
from icecream import IceCreamDebugger
|
|
import traceback
|
|
|
|
# To work with datatypes:
|
|
from typing import List
|
|
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** MACROS / ONE-TIME INIT ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
# --- Nothing Yet
|
|
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** VARIABLES ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
# --- Nothing Yet
|
|
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** CLASSES ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
class AsyncMySQL:
|
|
|
|
def __init__(
|
|
self,
|
|
pool_size,
|
|
*args,
|
|
**kwargs
|
|
):
|
|
|
|
"""
|
|
A class to work with SQL-based databases. Originally meant to only invoke stored procedures and retrieve them as
|
|
JSON-like structures (list or dict). The format for the results was very specific to our use case for serving
|
|
Bicree's requirement. This may not serve your requirement at all.
|
|
:param pool_size: The number of connections to maintain n a pool.
|
|
:param args: Any arguments to pass. Not used.
|
|
:param kwargs: Pass the connection configuration from here.
|
|
"""
|
|
|
|
# Set up the variables:
|
|
self.__args = args
|
|
self.__kwargs = kwargs
|
|
self.__min_pool_size = 10
|
|
self.__max_pool_size = max(pool_size, self.__min_pool_size)
|
|
self.__pool: aiomysql = None
|
|
|
|
# For establishing connection:
|
|
self.__exclusive_semaphore = asyncio.Semaphore(1)
|
|
|
|
# Minor adjustments for backward compatibility:
|
|
self.__kwargs["db"] = self.__kwargs.pop("database")
|
|
|
|
# Set up the debugging tools:
|
|
self.__printer = IceCreamDebugger(prefix = "MySQL | ", includeContext = True)
|
|
|
|
def __del__(self):
|
|
pass
|
|
|
|
async def connect(self) -> bool:
|
|
|
|
"""
|
|
Establish a connection and create a pool of connections to call from.
|
|
:return: True if connected, else False
|
|
"""
|
|
|
|
try:
|
|
|
|
# Ensure that only one connection attempt is being made at one time,
|
|
# and try to connect to the database:
|
|
async with self.__exclusive_semaphore:
|
|
if self.__pool is None:
|
|
self.__pool = await aiomysql.create_pool(
|
|
minsize = self.__min_pool_size,
|
|
maxsize = self.__max_pool_size,
|
|
loop = asyncio.get_event_loop(),
|
|
**self.__kwargs
|
|
)
|
|
|
|
# In case of any exception:
|
|
except Exception as exception:
|
|
self.__printer(exception)
|
|
self.__pool = None
|
|
|
|
# Done here:
|
|
return False if self.__pool is None else True
|
|
|
|
async def ensure_connection(self):
|
|
|
|
"""
|
|
Tries to ensure that a connection is present.
|
|
Can be called before every function to make sure that our pool is established.
|
|
:return: None.
|
|
"""
|
|
|
|
if self.__pool is None: await self.connect()
|
|
|
|
@staticmethod
|
|
def __parse_row(row):
|
|
|
|
"""
|
|
Converts from the custom objects of 'aiomysql' to types that are supported by Python.
|
|
:param row: The row from the result.
|
|
:return: The parsed row which will have types that are closer to being native to Python..
|
|
"""
|
|
|
|
parsed_row = []
|
|
for item in row:
|
|
if isinstance(item, decimal.Decimal): parsed_row.append(float(item))
|
|
else: parsed_row.append(item)
|
|
return parsed_row
|
|
|
|
async def fetch_all(self, cursor):
|
|
|
|
# Make a variable to hold all the result sets.
|
|
# Needed for when the procedure responds with many "tables":
|
|
all_result_sets = []
|
|
|
|
# Iterate over all result sets,
|
|
# and process them one-by-one:
|
|
while True:
|
|
|
|
# Process the current result set:
|
|
this_result_set = []
|
|
result = await cursor.fetchall()
|
|
if not cursor.description: break
|
|
columns = [desc[0] for desc in cursor.description]
|
|
for row in result: this_result_set.append(dict(zip(columns, self.__parse_row(row))))
|
|
all_result_sets.append(this_result_set)
|
|
|
|
# Move to the next set,
|
|
# or break out of the loop if all done:
|
|
if not await cursor.nextset(): break
|
|
|
|
# Done here:
|
|
return all_result_sets
|
|
|
|
async def call_procedure(
|
|
self,
|
|
procedure_name: str,
|
|
procedure_args: tuple,
|
|
commit: bool = True
|
|
):
|
|
|
|
"""
|
|
To call stored procedures and retrieve all the responses.
|
|
:param procedure_name: The name of the stored procedure that must be called.
|
|
:param procedure_args: The args to be sent to the stored procedure.
|
|
:param commit: Whether, or not, you would like to commit the execution.
|
|
:return: The raw result set as received from the database.
|
|
"""
|
|
|
|
# Make sure we have a connection:
|
|
await self.ensure_connection()
|
|
|
|
# Make a variable to hold all the result sets.
|
|
# Needed for when the procedure responds with many "tables":
|
|
all_result_sets = []
|
|
|
|
# Call the procedure and get the results:
|
|
async with self.__pool.acquire() as connection:
|
|
async with connection.cursor() as cursor:
|
|
await cursor.callproc(procedure_name, procedure_args)
|
|
all_result_sets = await self.fetch_all(cursor)
|
|
if commit: await connection.commit()
|
|
|
|
# Done here:
|
|
return all_result_sets
|
|
|
|
async def call_procedure_and_get_json(
|
|
self,
|
|
procedure_name,
|
|
procedure_args,
|
|
commit: bool = True,
|
|
retry_count = 1,
|
|
backoff_seconds = 0.5,
|
|
backoff_multiplier = 1.1,
|
|
return_exception = False
|
|
):
|
|
|
|
"""
|
|
The method to call when you need to call a stored procedure and retrieve the response as a JSON-like object.
|
|
This is custom formatting based on the structure created by Mr. bhushan Thakkar in late April (2024).
|
|
:param procedure_name: The name of the stored procedure that must be called.
|
|
:param procedure_args: The args to be sent to the stored procedure.
|
|
:param commit: Whether, or not, you would like to commit the execution.
|
|
:param retry_count: The max. number of times to try in case one or more attempts fail.
|
|
:param backoff_seconds: The time to wait before making the next attempt if the retry count is more than 1.
|
|
:param backoff_multiplier: The factor that dictates how much to modify the time delay by when waiting to retry.
|
|
:param return_exception: Whether, or not, you would like to return the exception object if something goes wrong.
|
|
:return: The formatted response and the exception (if asked for).
|
|
"""
|
|
|
|
# Note down the start time:
|
|
start_ts = time.time()
|
|
|
|
# Try to get the data from the database:
|
|
results = []
|
|
exception = None
|
|
for _ in range(retry_count):
|
|
try: results = await self.call_procedure(
|
|
procedure_name = procedure_name,
|
|
procedure_args = procedure_args,
|
|
commit = commit
|
|
)
|
|
except Exception as exc: exception = exc
|
|
if exception is None: break
|
|
await asyncio.sleep(backoff_seconds)
|
|
backoff_seconds = backoff_seconds * backoff_multiplier
|
|
|
|
# If the results are blank:
|
|
if len(results) == 0:
|
|
formatted_results = {
|
|
"status": 0,
|
|
"message": "Please contact admin (NE)" if exception is None else "Please contact admin (E)",
|
|
"seconds": time.time() - start_ts,
|
|
"data": {}
|
|
}
|
|
if return_exception: return formatted_results, exception
|
|
else: return formatted_results
|
|
|
|
# Extract the very basic success or failure indicators:
|
|
formatted_results = {
|
|
"status": results[0][0]["status"],
|
|
"message": results[0][0].get("message", "ok"),
|
|
"seconds": 0.0,
|
|
"data": {}
|
|
}
|
|
|
|
# Handle the remaining keys of the zeroth result set:
|
|
for key, value in results[0][0].items():
|
|
if key not in formatted_results.keys():
|
|
formatted_results["data"][key] = value
|
|
|
|
# Format
|
|
for index in range(len(results)):
|
|
if index > 0: formatted_results["data"][f"rs{index-1}"] = results[index]
|
|
|
|
# Note down the time taken:
|
|
formatted_results["seconds"] = time.time() - start_ts
|
|
|
|
# Done here:
|
|
if return_exception: return formatted_results, exception
|
|
else: return formatted_results
|
|
|
|
async def execute_one(
|
|
self,
|
|
query: str,
|
|
commit: bool = True,
|
|
return_exception: bool = False
|
|
):
|
|
|
|
"""
|
|
Runs one command / query in SQL.
|
|
:param query: The query / command to run.
|
|
:param commit: Whether, or not, you would like to commit the execution.
|
|
:param return_exception: Whether, or not, you would like to return the exception from this function.
|
|
:return: Either just the result or the result and the exception.
|
|
"""
|
|
|
|
# Make sure we have a connection:
|
|
await self.ensure_connection()
|
|
|
|
# Start by assuming failure:
|
|
rows_affected = None
|
|
results = None
|
|
excp = None
|
|
|
|
try:
|
|
|
|
# Get a connection and execute the command:
|
|
async with self.__pool.acquire() as connection:
|
|
async with connection.cursor() as cursor:
|
|
rows_affected = await cursor.execute(query)
|
|
results = await self.fetch_all(cursor)
|
|
if commit: await connection.commit()
|
|
|
|
# SQL-specific errors:
|
|
except aiomysql.MySQLError as exception:
|
|
self.__printer("SQL Exception", exception)
|
|
excp = exception
|
|
|
|
# Other errors:
|
|
except Exception as exception:
|
|
self.__printer("Other Exception", exception)
|
|
excp = exception
|
|
|
|
# Done here:
|
|
if return_exception: return rows_affected, results, excp
|
|
else: return rows_affected, results
|
|
|
|
async def execute_many(
|
|
self,
|
|
query: str,
|
|
data: List[tuple],
|
|
commit: bool = True,
|
|
return_exception: bool = False
|
|
):
|
|
|
|
"""
|
|
Runs many commands / queries in SQL.
|
|
Consider the following example:
|
|
QUERY: "INSERT INTO pincodeMaster (pincode, city, state) VALUES (%s, %s, %s);"
|
|
DATA: [
|
|
('110001', 'New Delhi', 'Delhi'),
|
|
('500001', 'Hyderabad', 'Telangana'),
|
|
('600001', 'Chennai', 'Tamil Nadu')
|
|
]
|
|
:param query: The query / command to run.
|
|
:param data: The data to substitute into the query string.
|
|
:param commit: Whether, or not, you would like to commit the execution.
|
|
:param return_exception: Whether, or not, you would like to return the exception from this function.
|
|
:return: Either just the result or the result and the exception.
|
|
"""
|
|
|
|
# Make sure we have a connection:
|
|
await self.ensure_connection()
|
|
|
|
# Start by assuming failure:
|
|
rows_affected = None
|
|
results = None
|
|
excp = None
|
|
|
|
try:
|
|
|
|
# Get a connection and execute the command:
|
|
async with self.__pool.acquire() as connection:
|
|
async with connection.cursor() as cursor:
|
|
rows_affected = await cursor.executemany(query, data)
|
|
results = await self.fetch_all(cursor)
|
|
if commit: await connection.commit()
|
|
|
|
# SQL-specific errors:
|
|
except aiomysql.MySQLError as exception:
|
|
self.__printer("SQL Exception", exception)
|
|
excp = exception
|
|
|
|
# Other errors:
|
|
except Exception as exception:
|
|
self.__printer("Other Exception", exception)
|
|
excp = exception
|
|
|
|
# Done here:
|
|
if return_exception: return rows_affected, results, excp
|
|
else: return rows_affected, results
|
|
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** FUNCTIONS ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
# --- Nothing Yet
|
|
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** MAIN PROGRAM ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
pass
|