def get_cred_config() -> Dict[str, str]:
if "CLOUD_SQL_CREDENTIALS_SECRET" in os.environ:
name = os.environ["CLOUD_SQL_CREDENTIALS_SECRET"]
client = secretmanager.SecretManagerServiceClient()
response = client.access_secret_version(request={"name": name})
logger.info("Credentials pulled from CLOUD_SQL_CREDENTIALS_SECRET")
return json.loads(response.payload.data.decode("UTF-8"))
# [END cloudrun_user_auth_secrets]
else:
logger.info(
"CLOUD_SQL_CREDENTIALS_SECRET env var not set. Defaulting to environment variables."
)
if "DB_USER" not in os.environ:
raise Exception("DB_USER needs to be set.")
if "DB_PASSWORD" not in os.environ:
raise Exception("DB_PASSWORD needs to be set.")
if "DB_NAME" not in os.environ:
raise Exception("DB_NAME needs to be set.")
if "CLOUD_SQL_CONNECTION_NAME" not in os.environ:
raise Exception("CLOUD_SQL_CONNECTION_NAME needs to be set.")
return {
"DB_USER": os.environ["DB_USER"],
"DB_PASSWORD": os.environ["DB_PASSWORD"],
"DB_NAME": os.environ["DB_NAME"],
"DB_HOST": os.environ.get("DB_HOST", None),
"CLOUD_SQL_CONNECTION_NAME": os.environ["CLOUD_SQL_CONNECTION_NAME"],
}
def init_tcp_connection_engine(
db_config: Dict[str, str]) -> sqlalchemy.engine.base.Engine:
creds = credentials.get_cred_config()
db_user = creds["DB_USER"]
db_pass = creds["DB_PASSWORD"]
db_name = creds["DB_NAME"]
db_host = creds["DB_HOST"]
# Extract host and port from db_host
host_args = db_host.split(":")
db_hostname, db_port = host_args[0], int(host_args[1])
pool = sqlalchemy.create_engine(
# Equivalent URL:
# postgres+pg8000://<db_user>:<db_pass>@<db_host>:<db_port>/<db_name>
sqlalchemy.engine.url.URL(
drivername="postgresql+pg8000",
username=db_user, # e.g. "my-database-user"
password=db_pass, # e.g. "my-database-password"
host=db_hostname, # e.g. "127.0.0.1"
port=db_port, # e.g. 5432
database=db_name, # e.g. "my-database-name"
),
**db_config,
)
pool.dialect.description_encoding = None
logger.info("Database engine initialised from tcp connection")
return pool
def init_unix_connection_engine(
db_config: Dict[str, str]) -> sqlalchemy.engine.base.Engine:
creds = credentials.get_cred_config()
db_user = creds["DB_USER"]
db_pass = creds["DB_PASSWORD"]
db_name = creds["DB_NAME"]
db_socket_dir = creds.get("DB_SOCKET_DIR", "/cloudsql")
cloud_sql_connection_name = creds["CLOUD_SQL_CONNECTION_NAME"]
pool = sqlalchemy.create_engine(
# Equivalent URL:
# postgres+pg8000://<db_user>:<db_pass>@/<db_name>
# ?unix_sock=<socket_path>/<cloud_sql_instance_name>/.s.PGSQL.5432
sqlalchemy.engine.url.URL(
drivername="postgresql+pg8000",
username=db_user, # e.g. "my-database-user"
password=db_pass, # e.g. "my-database-password"
database=db_name, # e.g. "my-database-name"
query={
"unix_sock":
"{}/{}/.s.PGSQL.5432".format(
db_socket_dir,
cloud_sql_connection_name # e.g. "/cloudsql"
) # i.e "<PROJECT-NAME>:<INSTANCE-REGION>:<INSTANCE-NAME>"
},
),
**db_config,
)
pool.dialect.description_encoding = None
logger.info("Database engine initialised from unix conection")
return pool
def init_connection_engine() -> Dict[str, int]:
if os.getenv("TRAMPOLINE_CI", None):
logger.info("Using NullPool for testing")
db_config = {"poolclass": NullPool}
else:
db_config = {
# Pool size is the maximum number of permanent connections to keep.
"pool_size": 5,
# Temporarily exceeds the set pool_size if no connections are available.
"max_overflow": 2,
# The total number of concurrent connections for your application will be
# a total of pool_size and max_overflow.
# SQLAlchemy automatically uses delays between failed connection attempts,
# but provides no arguments for configuration.
# 'pool_timeout' is the maximum number of seconds to wait when retrieving a
# new connection from the pool. After the specified amount of time, an
# exception will be thrown.
"pool_timeout": 30, # 30 seconds
# 'pool_recycle' is the maximum number of seconds a connection can persist.
# Connections that live longer than the specified amount of time will be
# reestablished
"pool_recycle": 1800, # 30 minutes
}
if os.environ.get("DB_HOST"):
return init_tcp_connection_engine(db_config)
else:
return init_unix_connection_engine(db_config)
def get_cred_config() -> Dict[str, str]:
secret = os.environ.get("CLOUD_SQL_CREDENTIALS_SECRET")
if secret:
return json.loads(secret)
# [END cloudrun_user_auth_secrets]
else:
logger.info(
"CLOUD_SQL_CREDENTIALS_SECRET env var not set. Defaulting to environment variables."
)
if "DB_USER" not in os.environ:
raise Exception("DB_USER needs to be set.")
if "DB_PASSWORD" not in os.environ:
raise Exception("DB_PASSWORD needs to be set.")
if "DB_NAME" not in os.environ:
raise Exception("DB_NAME needs to be set.")
if "CLOUD_SQL_CONNECTION_NAME" not in os.environ:
raise Exception("CLOUD_SQL_CONNECTION_NAME needs to be set.")
return {
"DB_USER": os.environ["DB_USER"],
"DB_PASSWORD": os.environ["DB_PASSWORD"],
"DB_NAME": os.environ["DB_NAME"],
"DB_HOST": os.environ.get("DB_HOST", None),
"CLOUD_SQL_CONNECTION_NAME":
os.environ["CLOUD_SQL_CONNECTION_NAME"],
}
def shutdown() -> None:
# Find all Sessions in memory and close them.
close_all_sessions()
logger.info("All sessions closed.")
# Each connection was released on execution, so just formally
# dispose of the db connection if it's been instantiated
if db:
db.dispose()
logger.info("Database connection disposed.")
def save_vote(team: str, uid: str, time_cast: datetime.datetime) -> None:
# Preparing a statement before hand can help protect against injections.
stmt = sqlalchemy.text("INSERT INTO pet_votes (time_cast, candidate, uid)"
" VALUES (:time_cast, :candidate, :uid)")
# Using a with statement ensures that the connection is always released
# back into the pool at the end of statement (even if an error occurs)
with db.connect() as conn:
conn.execute(stmt, time_cast=time_cast, candidate=team, uid=uid)
logger.info("Vote for %s saved.", team)
def create_tables() -> None:
# This is called before any request on the main app, ensuring the database has been setup
logger.info("Creating tables")
global db
db = init_connection_engine()
# Create pet_votes table if it doesn't already exist
with db.connect() as conn:
conn.execute("CREATE TABLE IF NOT EXISTS pet_votes"
"( vote_id SERIAL NOT NULL, "
"time_cast timestamp NOT NULL, "
"candidate VARCHAR(6) NOT NULL, "
"uid VARCHAR(128) NOT NULL, "
"PRIMARY KEY (vote_id)"
");")
def shutdown_handler(signal: int, frame: FrameType) -> None:
logger.info("Signal received, safely shutting down.")
database.shutdown()
middleware.logging_flush()
print("Exiting process.", flush=True)
sys.exit(0)
