forked from openedx-unsupported/configuration
-
Notifications
You must be signed in to change notification settings - Fork 2
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
chore: Added script to drop tables from DB
- Loading branch information
1 parent
6d9e80b
commit 9f70cd7
Showing
2 changed files
with
182 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,181 @@ | ||
""" | ||
Script to drop tables from an RDS MySQL database while handling foreign key dependencies. | ||
Usage: | ||
python drop_dop_tables.py --db-host=my-db-host --db-name=my-db | ||
Arguments: | ||
--db-host The RDS database host. | ||
--db-name The database name. | ||
--dry-run Enable dry run mode (no actual changes). | ||
Environment Variables: | ||
DB_USERNAME The RDS database username (set via environment variable). | ||
DB_PASSWORD The RDS database password (set via environment variable). | ||
Functionality: | ||
- Drops specific tables only if they have had no activity in the last 12 months. | ||
- Handles foreign key constraints before dropping dependent tables. | ||
- Ensures safe execution using retries for AWS service interactions. | ||
Example: | ||
export DB_USERNAME=admin | ||
export DB_PASSWORD=securepass | ||
python drop_dop_tables.py --db-host=mydb.amazonaws.com --db-name=mydatabase --dry-run | ||
""" | ||
|
||
import boto3 | ||
import click | ||
import backoff | ||
from botocore.exceptions import ClientError | ||
import pymysql | ||
import logging | ||
from datetime import datetime, timedelta | ||
|
||
|
||
MAX_TRIES = 5 | ||
|
||
TABLES_TO_DROP = [ | ||
"oauth2_provider_trustedclient", # FK reference to oauth2_client | ||
"third_party_auth_providerapipermissions", # FK reference to oauth2_client | ||
"oauth2_client", | ||
"oauth2_grant", | ||
"oauth2_accesstoken", | ||
"oauth2_refreshtoken", | ||
"oauth_provider_consumer", | ||
"oauth_provider_nonce", | ||
"oauth_provider_scope", | ||
"oauth_provider_token", | ||
] | ||
FK_DEPENDENCIES = { | ||
"third_party_auth_providerapipermissions": "oauth2_client", | ||
"oauth2_provider_trustedclient": "oauth2_client", | ||
} | ||
|
||
# Configure logging | ||
LOGGER = logging.getLogger(__name__) | ||
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') | ||
|
||
|
||
|
||
class EC2BotoWrapper: | ||
def __init__(self): | ||
self.client = boto3.client("ec2") | ||
|
||
@backoff.on_exception(backoff.expo, ClientError, max_tries=MAX_TRIES) | ||
def describe_regions(self): | ||
return self.client.describe_regions() | ||
|
||
|
||
class RDSBotoWrapper: | ||
def __init__(self, **kwargs): | ||
self.client = boto3.client("rds", **kwargs) | ||
|
||
@backoff.on_exception(backoff.expo, ClientError, max_tries=MAX_TRIES) | ||
def describe_db_instances(self): | ||
return self.client.describe_db_instances() | ||
|
||
|
||
def connect_to_db(db_host, db_user, db_password, db_name): | ||
""" Establish a connection to the RDS MySQL database """ | ||
logging.info("Connecting to the database...") | ||
return pymysql.connect( | ||
host=db_host, | ||
user=db_user, | ||
password=db_password, | ||
database=db_name, | ||
cursorclass=pymysql.cursors.DictCursor | ||
) | ||
|
||
|
||
def drop_foreign_key(connection, db_name, table_name, referenced_table, dry_run): | ||
last_activity = get_last_activity_date(connection, table_name) | ||
if last_activity: | ||
one_year_ago = datetime.now() - timedelta(days=365) | ||
if last_activity > one_year_ago: | ||
logging.info(f"Skipping {table_name}: Last activity was on {last_activity}") | ||
return | ||
|
||
query = f""" | ||
SELECT CONSTRAINT_NAME FROM INFORMATION_SCHEMA.KEY_COLUMN_USAGE | ||
WHERE TABLE_SCHEMA = '{db_name}' AND TABLE_NAME = '{table_name}' AND REFERENCED_TABLE_NAME = '{referenced_table}'; | ||
""" | ||
with connection.cursor() as cursor: | ||
cursor.execute(query) | ||
result = cursor.fetchone() | ||
if result: | ||
constraint_name = result["CONSTRAINT_NAME"] | ||
drop_query = f"ALTER TABLE {table_name} DROP FOREIGN KEY {constraint_name};" | ||
if dry_run: | ||
logging.info(f"[Dry Run] Would drop foreign key {constraint_name} from {table_name}.") | ||
else: | ||
cursor.execute(drop_query) | ||
connection.commit() | ||
logging.info(f"Dropped foreign key {constraint_name} from {table_name}.") | ||
|
||
|
||
|
||
def get_last_activity_date(connection, table_name): | ||
""" Retrieve the last activity date for a table """ | ||
query = f""" | ||
SELECT MAX(GREATEST( | ||
COALESCE(UPDATE_TIME, '1970-01-01 00:00:00'), | ||
COALESCE(CREATE_TIME, '1970-01-01 00:00:00') | ||
)) AS last_activity | ||
FROM information_schema.tables | ||
WHERE TABLE_NAME = '{table_name}'; | ||
""" | ||
with connection.cursor() as cursor: | ||
cursor.execute(query) | ||
result = cursor.fetchone() | ||
if result and result["last_activity"]: | ||
return datetime.strptime(str(result["last_activity"]), "%Y-%m-%d %H:%M:%S") | ||
return None # If no activity, return None | ||
|
||
|
||
def drop_table(connection, table_name, dry_run): | ||
last_activity = get_last_activity_date(connection, table_name) | ||
if last_activity: | ||
one_year_ago = datetime.now() - timedelta(days=365) | ||
if last_activity > one_year_ago: | ||
logging.info(f"Skipping {table_name}: Last activity was on {last_activity}") | ||
return | ||
|
||
logging.info(f"Dropping table {table_name}...") | ||
if dry_run: | ||
logging.info(f"[Dry Run] Would drop table {table_name}.") | ||
else: | ||
with connection.cursor() as cursor: | ||
cursor.execute(f"DROP TABLE IF EXISTS {table_name}") | ||
connection.commit() | ||
logging.info(f"Table {table_name} dropped.") | ||
|
||
|
||
@click.command() | ||
@click.option('--db-host', required=True, help="RDS DB host") | ||
@click.option('--db-user', envvar='DB_USERNAME', required=True, help="RDS DB user (can be set via environment variable DB_USERNAME)") | ||
@click.option('--db-password', envvar='DB_PASSWORD', required=True, help="RDS DB password (can be set via environment variable DB_PASSWORD)") | ||
@click.option('--db-name', required=True, help="RDS DB name") | ||
@click.option('--dry-run', is_flag=True, help="Enable dry run mode (no actual changes)") | ||
def drop_tables(db_host, db_user, db_password, db_name, dry_run): | ||
""" | ||
A script to drop tables from an RDS database while handling foreign key dependencies. | ||
Table names are read from the provided file. | ||
""" | ||
try: | ||
connection = connect_to_db(db_host, db_user, db_password, db_name) | ||
|
||
for table, referenced_table in FK_DEPENDENCIES.items(): | ||
drop_foreign_key(connection, db_name, table, referenced_table, dry_run) | ||
|
||
for table in TABLES_TO_DROP: | ||
drop_table(connection, table, dry_run) | ||
|
||
connection.close() | ||
logging.info("Database cleanup completed successfully.") | ||
except Exception as e: | ||
logging.error(f"An error occurred: {e}") | ||
|
||
|
||
if __name__ == '__main__': | ||
drop_tables() |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1 @@ | ||
../jenkins/requirements.txt |