initial commit
This commit is contained in:
+401
@@ -0,0 +1,401 @@
|
||||
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",
|
||||
# )
|
||||
Reference in New Issue
Block a user