""" 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 = None # 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): """ Establish a connection and create a pool of connections to call from. :return: None. """ try: self.__pool = await aiomysql.create_pool( minsize = self.__min_pool_size, maxsize = self.__max_pool_size, loop = asyncio.get_event_loop(), **self.__kwargs ) except Exception as exception: self.__printer(exception) self.__pool = None 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