Squashed 'utils_v2/' content from commit 0fdc939c
git-subtree-dir: utils_v2 git-subtree-split: 0fdc939cad796fa02eb9d1f817c53d28cf8c2364
This commit is contained in:
@@ -0,0 +1,451 @@
|
||||
"""
|
||||
|
||||
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 disconnect(self) -> bool:
|
||||
|
||||
try:
|
||||
|
||||
# Try to close the pool and wait for the process to complete:
|
||||
self.__pool.close()
|
||||
await self.__pool.wait_closed()
|
||||
|
||||
# In case of any exception:
|
||||
except Exception as exception:
|
||||
self.__printer(exception)
|
||||
|
||||
# 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 | List],
|
||||
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
|
||||
Reference in New Issue
Block a user