402 lines
13 KiB
Python
402 lines
13 KiB
Python
import dataclasses
|
|
import datetime
|
|
import functools
|
|
import logging
|
|
import time
|
|
import typing
|
|
import uuid
|
|
|
|
import boto3
|
|
|
|
# import mysql.connector.connection
|
|
import psycopg2
|
|
import pymysql
|
|
from psycopg2 import sql
|
|
import psycopg2.extensions
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class Credentials:
|
|
username: str
|
|
password: str
|
|
|
|
@classmethod
|
|
def new(cls):
|
|
username_format = "u{}" # must start with a letter
|
|
return cls(
|
|
username=username_format.format(uuid.uuid4().hex),
|
|
password=uuid.uuid4().hex,
|
|
) # make sure it starts with a letter
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class DbConnection:
|
|
endpoint: str
|
|
port: int
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class DbInstance(DbConnection):
|
|
cluster_id: str
|
|
arn: str
|
|
reader_endpoint: typing.Optional[str] = None
|
|
database_name: typing.Optional[str] = None
|
|
user: typing.Optional[Credentials] = None
|
|
master_user: typing.Optional[Credentials] = None
|
|
|
|
|
|
def retry(timeout: int, backoff_seconds: int = 10, bubble_errors: typing.List[typing.Type[Exception]] = None):
|
|
"""
|
|
Tries calling a function {timeout} seconds until it succeeds or gives up and throw an error
|
|
:param timeout: Timeout in seconds
|
|
:param backoff_seconds: Wait time between retries
|
|
:param bubble_errors: List of exceptions to bubble up
|
|
:return: Wrapped function
|
|
"""
|
|
bubble_errors = tuple(bubble_errors or [])
|
|
|
|
def wrapped(func: typing.Callable):
|
|
@functools.wraps(func)
|
|
def func_retrier(*args, **kwargs):
|
|
waited = 0
|
|
started_at = datetime.datetime.now()
|
|
while True:
|
|
try:
|
|
logger.debug(f"Calling {func}")
|
|
result = func(*args, **kwargs)
|
|
logger.debug(f"{func} completed in {(datetime.datetime.now() - started_at)}")
|
|
return result
|
|
except KeyboardInterrupt:
|
|
raise
|
|
except tuple(bubble_errors):
|
|
raise
|
|
except Exception as e:
|
|
doze = min(timeout - waited, backoff_seconds)
|
|
times = (timeout - waited) // backoff_seconds
|
|
logger.debug(f"{func} failed, will try {times} times in {doze} sec", exc_info=True)
|
|
time.sleep(doze)
|
|
waited += backoff_seconds
|
|
if waited >= timeout:
|
|
raise TimeoutError(f"Couldn't get a successful result from {func} in {timeout} seconds") from e
|
|
|
|
return func_retrier
|
|
|
|
return wrapped
|
|
|
|
|
|
def paginate_by_marker(
|
|
func: callable,
|
|
list_field: str,
|
|
) -> typing.Iterable:
|
|
"""
|
|
Simplifies paginating over AWS results with a simpler interface than Paginators.
|
|
:param func: A function that accepts a parameter named `Marker`
|
|
:param list_field: Result field to iterate on
|
|
"""
|
|
marker = ""
|
|
while True:
|
|
result = func(Marker=marker)
|
|
|
|
items = result.get(list_field, [])
|
|
yield from items
|
|
|
|
if "Marker" not in result:
|
|
break
|
|
|
|
marker = result["Marker"]
|
|
|
|
|
|
class PostgresqlService:
|
|
db_type = "Postgresql"
|
|
db_engine = "aurora-postgresql"
|
|
db_engine_version = "10.11"
|
|
default_master_database_name = "postgres"
|
|
cluster_id_format = "pg-{}" # must start with a letter
|
|
tenant_database_name_format = "db{}" # must start with a letter
|
|
|
|
def __init__(self, region: str):
|
|
self.rds = boto3.client("rds", region_name=region)
|
|
|
|
def create_db(
|
|
self,
|
|
total_instances: int,
|
|
db_subnet_group_name: str,
|
|
node_type: str,
|
|
security_group_id: str,
|
|
engine_version: str = None,
|
|
backup_retention_days: int = 30,
|
|
public: bool = False,
|
|
tags: dict = None,
|
|
):
|
|
if not engine_version:
|
|
engine_version = self.db_engine_version
|
|
|
|
cluster_id = self.cluster_id_format.format(uuid.uuid4().hex)
|
|
|
|
master_creds = Credentials.new()
|
|
|
|
logger.info(
|
|
f"Creating {self.db_type} cluster",
|
|
extra=dict(
|
|
cluster_id=cluster_id,
|
|
engine_version=engine_version,
|
|
db_subnet_group_name=db_subnet_group_name,
|
|
security_group_id=security_group_id,
|
|
),
|
|
)
|
|
|
|
db_result = self.rds.create_db_cluster(
|
|
DBClusterIdentifier=cluster_id,
|
|
Engine=self.db_engine,
|
|
EngineVersion=engine_version,
|
|
MasterUsername=master_creds.username,
|
|
MasterUserPassword=master_creds.password,
|
|
Tags=[{"Key": k, "Value": v} for k, v in (tags or {}).items()],
|
|
DBSubnetGroupName=db_subnet_group_name,
|
|
VpcSecurityGroupIds=[security_group_id],
|
|
BackupRetentionPeriod=backup_retention_days,
|
|
)["DBCluster"]
|
|
|
|
endpoint = db_result["Endpoint"]
|
|
reader_endpoint = db_result["ReaderEndpoint"]
|
|
port = db_result["Port"]
|
|
arn = db_result["DBClusterArn"]
|
|
|
|
for i in range(total_instances):
|
|
instance_id = f"{cluster_id}-Instance-{i}"
|
|
logger.info(
|
|
"Creating instances on the cluster",
|
|
extra=dict(
|
|
cluster_id=cluster_id,
|
|
instance_id=instance_id,
|
|
node_type=node_type,
|
|
),
|
|
)
|
|
instance_result = self.rds.create_db_instance(
|
|
DBInstanceIdentifier=instance_id,
|
|
DBClusterIdentifier=cluster_id,
|
|
DBInstanceClass=node_type,
|
|
Engine=self.db_engine,
|
|
PubliclyAccessible=public,
|
|
Tags=[{"Key": k, "Value": v} for k, v in (tags or {}).items()],
|
|
)
|
|
|
|
tenant_database_name = self.tenant_database_name_format.format(uuid.uuid4().hex)
|
|
db = DbInstance(
|
|
cluster_id=cluster_id,
|
|
arn=arn,
|
|
endpoint=endpoint,
|
|
reader_endpoint=reader_endpoint,
|
|
port=port,
|
|
database_name=tenant_database_name,
|
|
master_user=master_creds,
|
|
)
|
|
print(db) # TODO: remove
|
|
|
|
logger.info("Waiting until DB instances are online. This will take a while (~5 min)")
|
|
logger.info("Connecting DB instance using master credentials", extra=dict(endpoint=db.endpoint))
|
|
self._check_connection(
|
|
connection=db,
|
|
credentials=master_creds,
|
|
)
|
|
|
|
tenant_creds = Credentials.new()
|
|
logger.info(
|
|
"Connection successful. Creating a new database and user",
|
|
extra=dict(
|
|
endpoint=endpoint,
|
|
username=tenant_creds.username,
|
|
database_name=tenant_database_name,
|
|
),
|
|
)
|
|
self.create_tenant(
|
|
connection=db,
|
|
master_user=master_creds,
|
|
database_name=tenant_database_name,
|
|
tenant_user=tenant_creds,
|
|
)
|
|
db.user = tenant_creds
|
|
|
|
return db
|
|
|
|
@retry(timeout=10 * 60)
|
|
def _check_connection(self, connection: DbConnection, credentials: Credentials, database_name: str = "postgres"):
|
|
psycopg2.connect(
|
|
host=connection.endpoint,
|
|
port=connection.port,
|
|
dbname=database_name,
|
|
user=credentials.username,
|
|
password=credentials.password,
|
|
connect_timeout=15,
|
|
).close()
|
|
|
|
def _find_instance(self, endpoint: str) -> typing.Optional[DbInstance]:
|
|
func = functools.partial(
|
|
self.rds.describe_db_clusters, Filters=[{"Name": "engine", "Values": [self.db_engine]}], MaxRecords=100
|
|
)
|
|
cluster = None
|
|
for it in paginate_by_marker(func, "DBClusters"):
|
|
if it["Endpoint"] == endpoint:
|
|
cluster = it
|
|
break
|
|
if cluster:
|
|
return DbInstance(
|
|
endpoint=cluster["Endpoint"],
|
|
port=cluster["Port"],
|
|
cluster_id=cluster["DBClusterIdentifier"],
|
|
arn=cluster["DBClusterArn"],
|
|
reader_endpoint=cluster["ReaderEndpoint"],
|
|
)
|
|
return None
|
|
|
|
def create_tenant(
|
|
self,
|
|
connection: DbConnection,
|
|
master_user: Credentials,
|
|
tenant_user: Credentials,
|
|
database_name: typing.Optional[str] = None,
|
|
) -> DbInstance:
|
|
if not database_name:
|
|
database_name = self.tenant_database_name_format.format(uuid.uuid4().hex)
|
|
try:
|
|
con = psycopg2.connect(
|
|
host=connection.endpoint,
|
|
port=connection.port,
|
|
dbname=self.default_master_database_name,
|
|
user=master_user.username,
|
|
password=master_user.password,
|
|
connect_timeout=30,
|
|
)
|
|
con.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
|
|
cur = con.cursor()
|
|
|
|
cur.execute(
|
|
sql.SQL("CREATE DATABASE {}").format(
|
|
sql.Identifier(database_name),
|
|
),
|
|
)
|
|
cur.execute(
|
|
sql.SQL("CREATE USER {} WITH ENCRYPTED PASSWORD {}").format(
|
|
sql.Identifier(tenant_user.username),
|
|
sql.Placeholder(),
|
|
),
|
|
[tenant_user.password],
|
|
)
|
|
cur.execute(
|
|
sql.SQL("GRANT ALL PRIVILEGES ON DATABASE {} TO {}").format(
|
|
sql.Identifier(database_name),
|
|
sql.Identifier(tenant_user.username),
|
|
)
|
|
)
|
|
|
|
# `with` block would create an implicit transaction
|
|
con.close()
|
|
|
|
db = self._find_instance(connection.endpoint)
|
|
db.database_name = database_name
|
|
db.user = tenant_user
|
|
db.master_user = master_user
|
|
return db
|
|
except psycopg2.Error:
|
|
logging.error(
|
|
f"Failed to create database and user",
|
|
exc_info=True,
|
|
extra=dict(endpoint=connection.endpoint, database_name=database_name),
|
|
)
|
|
raise
|
|
|
|
|
|
class MysqlService(PostgresqlService):
|
|
db_type = "MySQL"
|
|
db_engine = "aurora-mysql"
|
|
db_engine_version = "5.7.12"
|
|
default_master_database_name = "mysql"
|
|
cluster_id_format = "mysql-{}" # must start with a letter
|
|
|
|
@retry(timeout=10 * 60)
|
|
def _check_connection(self, connection: DbConnection, credentials: Credentials, database_name: str = "postgres"):
|
|
pymysql.connect(
|
|
host=connection.endpoint,
|
|
port=connection.port,
|
|
user=credentials.username,
|
|
password=credentials.password,
|
|
connect_timeout=15,
|
|
).close()
|
|
|
|
def create_tenant(
|
|
self,
|
|
connection: DbConnection,
|
|
master_user: Credentials,
|
|
tenant_user: Credentials,
|
|
database_name: typing.Optional[str] = None,
|
|
) -> DbInstance:
|
|
if not database_name:
|
|
database_name = self.tenant_database_name_format.format(uuid.uuid4().hex)
|
|
|
|
try:
|
|
with pymysql.connect(
|
|
host=connection.endpoint,
|
|
port=connection.port,
|
|
user=master_user.username,
|
|
password=master_user.password,
|
|
connect_timeout=15,
|
|
) as con:
|
|
cur = con.cursor()
|
|
cur.execute(f"CREATE DATABASE {database_name}")
|
|
cur.execute(f"CREATE USER {tenant_user.username} IDENTIFIED BY '{tenant_user.password}'")
|
|
cur.execute(
|
|
f"GRANT ALL PRIVILEGES ON {database_name}.* TO {tenant_user.username}@'%' IDENTIFIED BY '{tenant_user.password}'"
|
|
)
|
|
|
|
db = self._find_instance(endpoint=connection.endpoint)
|
|
db.database_name = database_name
|
|
db.master_user = master_user
|
|
db.user = tenant_user
|
|
return db
|
|
except pymysql.Error:
|
|
logger.error(
|
|
"Failed to create database and user",
|
|
exc_info=True,
|
|
extra=dict(endpoint=connection.endpoint, database_name=database_name),
|
|
)
|
|
raise
|
|
|
|
|
|
if __name__ == "__main__":
|
|
logging.basicConfig(level=logging.INFO, format=f"%(asctime)s: {logging.BASIC_FORMAT}")
|
|
logging.getLogger("botocore").setLevel(logging.INFO)
|
|
|
|
# m = MysqlService("eu-central-1")
|
|
# db = m.create_db(
|
|
# total_instances=1,
|
|
# db_subnet_group_name="mydbsubnet",
|
|
# security_group_id="sg-c11a77ac",
|
|
# node_type="db.t3.medium",
|
|
# tags={"CostCenter": "hello"},
|
|
# public=True,
|
|
# )
|
|
# print(db)
|
|
|
|
# p = PostgresqlService("eu-central-1")
|
|
# db = p.create_db(
|
|
# total_instances=1,
|
|
# db_subnet_group_name="mydbsubnet",
|
|
# security_group_id="sg-c11a77ac",
|
|
# node_type="db.t3.medium",
|
|
# tags={"CostCenter": "hello"},
|
|
# public=True,
|
|
# )
|
|
# p.create_tenant(
|
|
# connection=DbConnection(
|
|
# endpoint="pg0a085614d5eb4f34affd2a24b1fb18c4.cluster-c8v5rp0ouaey.eu-central-1.rds.amazonaws.com",
|
|
# port=5432,
|
|
# ),
|
|
# master_user=Credentials(
|
|
# username="u7d0f9621f5a5419288a4ef7ce9fd5fe", password="ca9473879b794dd3a0b0f0d2a8d4aebb"
|
|
# ),
|
|
# tenant_user=Credentials.new(),
|
|
# database_name="mydb",
|
|
# )
|