import uuid
from airflow.providers.postgres.operators.postgres import PostgresOperator
from airflow.providers.postgres.hooks.postgres import PostgresHook
from airflow.exceptions import AirflowFailException
from config import Config
from util.s3 import S3Util
from util.ecs import ECSUtil
from util.dag import DagUtil
UNLOAD_QUERY = """
UNLOAD (
{tgt_query}
)
TO '{tgt_pth}'
WITH (
format = 'parquet',
compression='snappy'
)
"""
GRAPI_PROD_VAL = "PROD grapi"
GRAPI_STAGE_VAL = "STAGE grapi"
GRAPI_DEV_VAL = "DEV grapi"
[docs]
class DBUtil:
"""
Central database utilities
"""
[docs]
@staticmethod
def add_series_column(
conn_id: str,
table: str,
column: str = "id",
):
"""
Adds a series column to a table
:param conn_id: the connection id
:param table: the table to add the column to
:param column: the column name
:return: the postgres operator that adds the column
"""
return PostgresOperator(
task_id=f"create_series_{conn_id}_{table}_{column}",
postgres_conn_id=conn_id,
autocommit=True,
sql=f"ALTER TABLE {table} ADD COLUMN IF NOT EXISTS {column} SERIAL PRIMARY KEY;",
)
[docs]
@staticmethod
def reset_sequence(
conn_id: str,
table: str,
column: str = "id",
seq_name: str = None,
):
"""
Resets the sequence for a table
"""
seq_name = seq_name or f"{table}_{column}_seq"
return PostgresOperator(
task_id=f"reset_{table}_{column}_seq",
postgres_conn_id=conn_id,
autocommit=True,
sql=f"""
select setval('{seq_name}', (select max({column})+1 from {table}), false);
""",
)
[docs]
@staticmethod
def table_exists(
db_hook,
schema: str,
table: str,
) -> bool:
"""
Checks if a postgres table exists
:param db_hook: the postgres hook
:param schema: the schema to check
:param table: the table to check
:return: True if the table exists
"""
result = db_hook.get_records(f"SELECT to_regclass('{schema}.{table}');")
return result and result[0] and result[0][0]
[docs]
@staticmethod
def schema_exists(
db_hook,
schema: str,
) -> bool:
"""
Checks if a postgres schema exists
:param db_hook: the postgres hook
:param schema: the schema to check
:return: True if the schema exists
"""
result = db_hook.get_records(
f"SELECT EXISTS (SELECT 1 FROM information_schema.schemata WHERE schema_name = '{schema}');"
)
return result and result[0] and result[0][0]
[docs]
@staticmethod
def generate_case_from_dict(
tgt_field: str,
case_dict: dict,
) -> str:
"""
Generates a case statement from a dictionary
:param tgt_field: the target field
:param case_dict: the dictionary to use
:return: the case statement
"""
return "".join(
[
f"when {tgt_field} = '{key}' then '{value}'\n"
for key, value in case_dict.items()
]
)
[docs]
@staticmethod
def upsert_df(
df,
table_name,
engine,
schema_name="public",
) -> bool:
"""
Implements the equivalent of pd.DataFrame.to_sql(..., if_exists='update')
(which does not exist). Creates or updates the db records based on the
dataframe records.
Conflicts to determine update are based on the dataframes index.
This will set unique keys constraint on the table equal to the index names
1. Create a temp table from the dataframe
2. Insert/update from temp table into table_name
Returns: True if successful
:param df: the dataframe to upsert
:type df: pd.DataFrame
:param table_name: the table name
:type table_name: str
:param engine: the sqlalchemy engine
:type engine: sqlalchemy.engine.base.Connection
:param schema_name: the schema name
:type schema_name: str
:return: True if successful
"""
if len(df) == 0:
print("Func upsert_df is called on empty df >> Return False/failed")
return False
# If the table does not exist, we should just use to_sql to create it
if not engine.execute(
f"""SELECT EXISTS (
SELECT FROM information_schema.tables
WHERE table_schema = '{schema_name}'
AND table_name = '{table_name}');
"""
).first()[0]:
df.to_sql(table_name, engine, schema=schema_name)
return True
# If it already exists...
temp_table_name = f"temp_{uuid.uuid4().hex[:6]}"
df.to_sql(temp_table_name, engine, index=True, schema=schema_name)
index = list(df.index.names)
colnames_index_pkey = ", ".join([f'"{i}"' for i in index])
columns = list(df.columns)
headers = index + columns
headers_sql_txt = ", ".join(
[f'"{i}"' for i in headers]
) # index1, index2, ..., column 1, col2, ...
# col1 = exluded.col1, col2=excluded.col2
update_column_stmt = ", ".join(
[f'"{col}" = EXCLUDED."{col}"' for col in columns]
)
# For the ON CONFLICT clause, postgres requires that the columns have unique constraint
# # This "DROP CONSTRAINT IF EXISTS" fails to identify existing UNIQUE CONSTRAINT with a different name
# query_pk = f"""
# ALTER TABLE "{table_name}" DROP CONSTRAINT IF EXISTS unique_constraint_for_upsert_{table_name};
# ALTER TABLE "{table_name}" ADD CONSTRAINT unique_constraint_for_upsert_{table_name} UNIQUE ({colnames_index_pkey});
# """
# engine.execute(query_pk)
## ALTERNATIVELY, CREATE UNIQUE INDEX: https://stackoverflow.com/a/74057278; https://www.postgresql.org/docs/current/sql-insert.html
pg_sql_unique_constraint_index = f"""
CREATE UNIQUE INDEX IF NOT EXISTS unique_constraint_for_upsert_{schema_name}_{table_name} ON {schema_name}.{table_name}({colnames_index_pkey})
"""
engine.execute(pg_sql_unique_constraint_index)
# Compose and execute upsert query. Ref: https://www.postgresqltutorial.com/postgresql-tutorial/postgresql-upsert/
query_upsert = f"""
INSERT INTO {schema_name}.{table_name} ({headers_sql_txt})
SELECT {headers_sql_txt} FROM {schema_name}.{temp_table_name}
ON CONFLICT ({colnames_index_pkey}) DO UPDATE
SET {update_column_stmt};
"""
engine.execute(query_upsert)
engine.execute(f"DROP TABLE {schema_name}.{temp_table_name}")
return True
[docs]
@staticmethod
def transfer_datalake_db(
temp_key: str,
table_name: str,
tgt_query: str,
tgt_bucket: str = Config.DATA_LAKE_BUCKET,
data_lake_db: str = Config.DATA_LAKE_DB,
workgroup: str = Config.DEFAULT_WORKGROUP,
batch_size: int = Config.BATCH_SIZE_ECS,
dest_db: str = Config.READY_DATA_DB,
create_suffix: str = "",
drop_table: bool = False,
use_pg_copy: bool = False,
task_suffix: str = "",
on_conflict_sql: str | None = None,
upsert_origin_table: str | None = None,
):
"""
Transfers from the data lake to a postgres db
Sequentially loads parquet files from s3 to postgres
This probably could be an operator
:param temp_key: the temp key to store data
:type temp_key: str
:param table_name: the table name in the destination db
:type table_name: str
:param tgt_query: the target query in athena
:type tgt_query: str
:param tgt_bucket: the target bucket
:type tgt_bucket: str
:param data_lake_db: the data lake db
:type data_lake_db: str
:param workgroup: the workgroup
:type workgroup: str
:param batch_size: the batch size
:type batch_size: int
:param dest_db: the destination db
:type dest_db: str
:param create_suffix: if drop_table, something to add at the end of the create table query
:type create_suffix: str
:param drop_table: whether to drop the table before query runs
:type drop_table: bool
:param use_pg_copy: whether to use pg copy (usually faster and more efficient)
:type use_pg_copy: bool
:param task_suffix: the task suffix, large or small
:type task_suffix: str
:return: task arn
:rtype: _type_
"""
import awswrangler as wr
S3Util.delete_s3_folder(tgt_bucket, temp_key)
query = UNLOAD_QUERY.format(
tgt_query=tgt_query,
tgt_pth=f"s3://{tgt_bucket}/{temp_key}",
)
print(query)
wr.athena.start_query_execution(
sql=query,
database=data_lake_db,
workgroup=workgroup,
wait=True,
)
ECSUtil.run_s3_parquet_to_pg(
table=table_name,
key=temp_key,
drop_table=drop_table,
task_suffix=task_suffix,
conn_id=dest_db,
batch_size=batch_size,
suffix=create_suffix,
use_pg_copy=use_pg_copy,
on_conflict_sql=on_conflict_sql,
upsert_origin_table=upsert_origin_table,
)
[docs]
@staticmethod
def get_pg_columns(
table_name: str,
schema: str = "public",
database_name: str = Config.READY_DATA_DB,
) -> list[str]:
"""
Returns an array of the columns of a table in pg
:param table_name: name of the table
:type table_name: str
:param schema: table schema
:type schema: str
:param database_name: database name
:type database_name: str
:return: an array of strings or empty array if table does not exist
:rtype: _type_
"""
db_hook = PostgresHook(postgres_conn_id=database_name)
if not DBUtil.table_exists(db_hook=db_hook, schema=schema, table=table_name):
return []
sql = f"""SELECT attrelid::regclass AS tbl
, attname AS col
, atttypid::regtype AS datatype
-- more attributes?
FROM pg_attribute
WHERE attrelid = '{schema}.{table_name}'::regclass -- table name optionally schema-qualified
AND attnum > 0
AND NOT attisdropped
ORDER BY attnum;"""
records = db_hook.get_records(sql)
return [row[1] for row in records]
[docs]
@staticmethod
def pg_swap_table(
db_hook,
stage_table: str,
table: str,
postfix: str = "_old",
) -> bool:
"""
Within the public schema, swaps a stage table name for a production table, then drops the old table.
Please note that this is a destructive operation.
It is assumed that indexes and data have been transferred by the time this is called.
:param db_hook: the postgres hook
:type db_hook: _type_
:param stage_table: the name of the stage table that will become production
:type stage_table: str
:param table: the name of the production table
:type table: str
:param postfix: the postfix to add to the production table that was replaced
:type postfix: str
:return: True if both queries execute successfully
"""
old_table = f"{table}{postfix}"
try:
db_hook.run(
f"ALTER TABLE {table} RENAME TO {old_table};",
autocommit=True,
)
db_hook.run(
f"ALTER TABLE {stage_table} RENAME TO {table};",
autocommit=True,
)
except AttributeError as e:
print(e)
raise AirflowFailException("Error renaming table")
else:
if DBUtil.table_exists(db_hook, "public", table):
try:
db_hook.run(f"DROP TABLE {old_table};", autocommit=True)
except AttributeError as e:
print(e)
raise AirflowFailException(f"Error dropping old table {old_table}")
else:
print(
f"!!Table swap failed. Production table {table} does not exist after swap!!"
)
print("Attempting to rename back to original state")
if DBUtil.table_exists(db_hook, "public", old_table):
print(
f"Original table {table} had been renamed to {old_table}; hotswap failed >> now proceeding with revert"
)
try:
db_hook.run(
f"ALTER TABLE {old_table} RENAME TO {table};",
autocommit=True,
)
except AttributeError as e:
print(e)
raise AirflowFailException(
"!!Revert has failed, please check the database!!"
)
else:
print("Revert successful")
raise AirflowFailException(
"!!Table swap has failed, but the production table has been successfully reverted!!"
)
print("Table swap successful")
return True
[docs]
@staticmethod
def pg_table_partitions(
db_hook,
schema: str,
table_name: str,
) -> list[str]:
"""
gets a list of the partitions of a table. returns empty list if table does not exist
:param db_hook: pg db hook
:type db_hook: _type_
:param schema: table schema
:type schema: str
:param table_name: table name
:type table_name: str
:return: an array of the table names
:rtype: _type_
"""
is_src_table_available = DBUtil.table_exists(db_hook, schema, table_name)
if is_src_table_available:
print(f"{table_name} does exist")
# get existing partitions/ probably fails if table does not exist
result_partitions = db_hook.get_records(
f"""SELECT CASE
WHEN POSITION('.' IN inhrelid::regclass::text) > 0
THEN substring(inhrelid::regclass::text from '.*\.(.*)')
ELSE inhrelid::regclass::text
END AS child
FROM pg_catalog.pg_inherits
WHERE inhparent = '{schema}.{table_name}'::regclass; """
) # case stmt to remove schema name from the partition name (Pg Version 16+; else keep as is)
return [partition[0] for partition in result_partitions]
return []
[docs]
@staticmethod
def pg_table_indexes(
db_hook,
schema: str,
table_name: str,
):
"""
gets a list of indexes on a table. it Includes all the indexes in the partitions too
:param db_hook: _description_
:type db_hook: _type_
:param schema: _description_
:type schema: str
:param table_name: _description_
:type table_name: str
:param suffix: _description_, defaults to '_old'
:type suffix: str, optional
"""
result_partitions = DBUtil.pg_table_partitions(db_hook, schema, table_name)
result_partitions.append(table_name)
table_indexes = []
for table in result_partitions:
all_indexes_sql = (
f"""SELECT indexname FROM pg_indexes WHERE tablename='{table}';"""
)
list_indexes = db_hook.get_records(all_indexes_sql)
for row in list_indexes:
table_indexes.append(row[0])
return table_indexes
[docs]
@staticmethod
def boss_db_param_values() -> dict[str, str]:
"""
list BOSS DB parameters
:return: the dict of BOSS DB parameters
:rtype: dict[str, str]
"""
dbs = (
{
"BOSS_PLATFORM_DB": "Boss Platform Prod",
"BOSS_PLATFORM_STAGE_DB": "Boss Platform Stage",
}
if DagUtil.is_prod_instance()
else {
"BOSS_PLATFORM_DB": "Boss Platform Dev",
}
)
return dbs
[docs]
@staticmethod
def pg_swap_with_partitions(
db_hook,
stage_table: str,
table: str,
postfix: str = "_old",
schema: str = "public",
drop_backup_at_end: bool = False,
drop_backup_before_start=False,
):
"""
This is similar to the pg_swap_table but with
experimental function to support swapping tables with partitions
Highly Experimental. better to use pg_swap_table
TODO: After enough testing combine these two functions into one
This should be an util. we already have one util similar to this but need to add partition name support to it
for now I kept it here to test it more and make it work with tables without partitions too and combine it with the existing util
:param db_hook: the postgres hook
:param stage_table: the name of the stage table that will become production. it Must have _stage at the end of its table name and partitions
:param table: the name of the production table
:param postfix: the postfix to add to the production table that was replaced
:param drop_backup_at_end: it will drop the backup table at the end of the process
:param drop_backup_before_start: it will drop {table}_{postfix} table before start the process
:return: True if both queries execute successfully
"""
def rename_table_partitions(
table_name, schema, postfix, db_hook, remove_postfix=False
):
"""
This will be an internal function in db_util. right now it is kept here
"""
result_partitions = DBUtil.pg_table_partitions(db_hook, schema, table_name)
# check if we have dest partitions already. if yes fail. old_name:new_name
new_partition_names = {
f"{partition}": (
f"{partition}".replace(postfix, "")
if remove_postfix
else f"{partition}{postfix}"
)
for partition in result_partitions
}
# src table itself should also be checked
new_table_name = (
f"{table_name}".replace(postfix, "")
if remove_postfix
else f"{table_name}{postfix}"
)
new_partition_names[table_name] = new_table_name
print(new_partition_names)
# rename tables
sql_queries = []
for partition in new_partition_names:
if partition == new_partition_names[partition]:
print(
f"{partition} is supposed to change to {new_partition_names[partition]}. Skipping"
)
continue
old_name = partition.replace(f"{schema}.", "")
new_name = new_partition_names[partition].replace(f"{schema}.", "")
sql_queries.append(
f"ALTER TABLE IF EXISTS {old_name} RENAME TO {new_name};"
)
return sql_queries
try:
sql_commands = []
stage_postfix = "_stage"
if not stage_table.endswith(stage_postfix):
raise AirflowFailException(f"stage_table must end with {stage_postfix}")
if drop_backup_before_start:
sql_commands.append(
f"DROP TABLE IF EXISTS {schema}.{table}{postfix} CASCADE;"
)
# rename prod table with its partitions to {prod_origin_name}_{postfix}
if DBUtil.table_exists(db_hook=db_hook, schema=schema, table=table):
rename_old = rename_table_partitions(
table_name=table,
schema=schema,
postfix=postfix,
remove_postfix=False,
db_hook=db_hook,
)
sql_commands.extend(rename_old)
# rename stage table to prod table {prod_origin_name}
rename_new = rename_table_partitions(
table_name=stage_table,
schema=schema,
postfix=stage_postfix,
remove_postfix=True,
db_hook=db_hook,
)
sql_commands.extend(rename_new)
full_sql = f"""
begin;
SET search_path TO {schema};
{''.join(sql_commands)}
commit;
"""
print(f"{full_sql=}")
db_hook.run(
full_sql,
autocommit=True,
)
if drop_backup_at_end:
db_hook.run(
f"DROP TABLE IF EXISTS {schema}.{table}{postfix} CASCADE;",
autocommit=True,
)
except Exception as e:
raise Exception(
f"renaming partitions failed. {schema}.{table} still exist. {e}"
)
[docs]
@staticmethod
def get_game_db_connections():
"""
Get GAME database connections
"""
return (
{
Config.GAME_PG_STAGE: "GAME stage",
Config.GAME_PG: "GAME prod",
}
if Config.PRODUCTION
else {
Config.GAME_PG: "GAME dev",
}
)
[docs]
@staticmethod
def get_all_pg_conn_values():
"""
Get all the pg connections to use in the dag params as values_display
"""
return (
{
Config.BOSS_PLATFORM_GRAPI_DEV_DB: "GrAPI DEV",
Config.GAME_PG_DEV: "GAME DEV",
Config.BOSS_PLATFORM_GRAPI_STAGE_DB: "GrAPI STAGE",
Config.GAME_PG_STAGE: "GAME STAGE",
Config.BOSS_PLATFORM_GRAPI_DB: "GrAPI PROD",
Config.GAME_PG: "GAME PROD",
}
if Config.PRODUCTION
else {
Config.BOSS_PLATFORM_GRAPI_DB: "GrAPI DEV",
Config.GAME_PG: "GAME DEV",
}
)
[docs]
@staticmethod
def get_tileserver_environments():
"""
Get all the tileserver environments. Used in the dag params
:example:
```python
env_list = DBUtil.get_tileserver_environments()
Param(
default=list(env_list)[0],
type="string",
title="Destination Tileserver",
section="Deployment params",
description="Select Destination environment from a list",
enum=list(env_list.keys()),
values_display=env_list,
)
```
"""
if Config.PRODUCTION:
return {
"main": "Prod",
"staging": "Staging",
}
return {
"main": "Dev",
# "staging": "Staging", # not to be accessed on DEV
}
[docs]
@staticmethod
def get_athena_cross_dbs():
"""
Get all the Athena cross database list. Used in the dag params
"""
if Config.PRODUCTION:
return {
"ATHENA_GRAPI_DB_ALIAS": "Grapi Prod Via Athena",
"ATHENA_GRAPI_DB_ALIAS_STAGE": "Grapi Stage Via Athena",
"GAME_ATHENA": "GAME Prod Via Athena",
"GAME_ATHENA_STAGE": "GAME Stage Via Athena",
}
return {
"ATHENA_GRAPI_DB_ALIAS": "Grapi Dev Via Athena",
"ATHENA_GRAPI_DB_ALIAS_STAGE": "Grapi Dev Via Athena", # not to be accessed on DEV
"GAME_ATHENA_DEV": "GAME Dev Via Athena",
}
[docs]
@staticmethod
def get_athena_grapi_pg_aliases():
"""
Get all the Athena grapi PG aliases.
"""
return (
{
GRAPI_PROD_VAL: Config.GRAPI_ATHENA,
GRAPI_STAGE_VAL: Config.GRAPI_ATHENA_STAGE,
}
if Config.PRODUCTION
else {
GRAPI_DEV_VAL: Config.GRAPI_ATHENA,
}
)
[docs]
@staticmethod
def get_tenant_schemas(
conn_id: str = Config.BOSS_PLATFORM_GRAPI_DB,
include_public: bool = False,
) -> list[tuple[str]]:
"""
!WARNING!: This function makes a database call. Do not use it outside of a single task context. If you call this when constructing your DAG structure or params, the schedule will run this query every minute.
Query grapi for all tenant schema names, excludes public by default as that maps to the Ready tenant
:param conn_id: the connection id
:param include_public: whether to include the public schema
:return: the list of tenant schema names
"""
query = f"""
SELECT nspname
FROM pg_catalog.pg_namespace
where 1=1
and nspname not ilike 'pg_%'
and nspname not in (
'data',
'information_schema'{"" if include_public else ", 'public'"}
)
"""
db_hook = PostgresHook(postgres_conn_id=conn_id)
return db_hook.get_records(query)
# Could this be handled by BeadUtil? - MvE
[docs]
@staticmethod
def tenant_to_statefips(
conn_id: str = Config.BOSS_PLATFORM_GRAPI_DB,
) -> dict[str, str]:
"""
!WARNING!: This function makes a database call. Do not use it outside of a single task context. If you call this when constructing your DAG structure or params, the schedule will run this query every minute.
returns a dictionary of tenant schema name to state fips
:return: The tenant name and schema mapping dict
"""
tenants = DBUtil.get_tenant_schemas(conn_id)
tenants_list = [tenant[0] for tenant in tenants]
tenants_dict = {
Config.STATE_NAMES[state_abbr]
.lower()
.replace(" ", "_"): Config.STATE_FIPS[state_abbr]
for state_abbr in Config.STATE_FIPS.keys()
if state_abbr in Config.STATE_NAMES.keys()
}
tenants_dict["alabama_scabcd"] = Config.STATE_FIPS["AL"]
tenants_dict["michigan_detroit"] = Config.STATE_FIPS["MI"]
return {k: v for k, v in tenants_dict.items() if k in tenants_list}
[docs]
@staticmethod
def run_env_db(
query: str,
prod_db: str = Config.BOSS_PLATFORM_GRAPI_DB,
stage_db: str = Config.BOSS_PLATFORM_GRAPI_STAGE_DB,
):
"""
Helper to run queries against stage + prod or just dev
:param query: the query to run
:param prod_db: the prod db
:param stage_db: the stage db (only executes when running on prod)
"""
prod_hook = PostgresHook(postgres_conn_id=prod_db)
prod_hook.run(query)
if DagUtil.is_prod_instance():
stage_hook = PostgresHook(postgres_conn_id=stage_db)
stage_hook.run(query)