from typing import List, Optional, Union, Dict import pandas as pd from internal.helpers import get_rds_connection, make_iso_date_now from internal.queries import rds_query from internal.commons import EFS_PATH from .aws_connectors import run_query, pd_read_sql from io import StringIO from utils.async_task_manager import io_task from utils import appsync_communication as appsync import logging component_logger = logging.getLogger() VISIBLE_COLLECTIONS = ['source', 'collection set', 'segment'] def add_event_log(schema: str, collection_id: int, user_id: str, log_type: str, description: str): run_query(f"""INSERT INTO {schema}.event_log (collection_id, user_id, type, description) VALUES (%(collection_id)s, %(user_id)s, %(log_type)s,%(description)s);""", {'collection_id': collection_id, 'user_id': user_id, 'log_type': log_type, 'description': description}, fetch=False) def generate_collection(collection_name: str, collection_source: str, schema_name: str, user_id: str, parent_id: int = None, status: str = 'undefined', collection_type: str = 'source', log_description: str = 'not specified') -> int: """ Loads collection into collection table and return collection id""" sql = f""" INSERT INTO {schema_name}.collection (name, source, parent_id, status, collection_type) VALUES (%(collection_name)s, %(collection_source)s, %(parent_id)s, %(status)s, %(collection_type)s ) RETURNING id;""" collection_id = run_query(sql, {'collection_name': collection_name, 'collection_source': collection_source, 'parent_id': parent_id, 'status': status, 'collection_type': collection_type })[0][0] """ Update event log """ add_event_log(schema=schema_name, collection_id=collection_id, user_id=user_id, log_type='collection_creation', # TODO! this has to be from some agreed enum description=log_description) component_logger.info(f"Generated collection {collection_id}:'{collection_name}' dataframe") change_collection_status(schema_name, user_id, collection_id, status) return collection_id def get_collection_type(schema: str, collection_id: int) -> Optional[str]: """ Gets collection type """ for collection_type in run_query(f"SELECT collection_type FROM {schema}.collection WHERE id = %(collection_id)s", {'collection_id': collection_id}): component_logger.info(f'Got: {collection_type}') return collection_type[0] return None def get_management_schema(schema: str, for_fan_table=False) -> str: """ Returns management schema for alliances/workspaces. If it's company schema, just returns itself.""" if for_fan_table and schema.startswith('a'): return schema if schema.startswith('c'): return schema for mgm in run_query(f"SELECT management_company_id FROM commons.company_alliance WHERE alliance_id = %(schema)s", {'schema': schema}): return mgm[0] raise RuntimeError('Unknown schema') def get_collection_attribute_names(schema, collection_id=None): """ Gets attribute names used by collection """ sql = f"""SELECT a.name FROM {schema}.attribute a """ if collection_id is not None: sql += f"INNER JOIN {schema}.collection_attribute ca ON a.id = ca.attribute_id " \ f"INNER JOIN {schema}.collection c ON ca.collection_id = c.id " \ f"WHERE c.id = %(collection_id)s OR c.parent_id = %(collection_id)s;" attribute_names = run_query(sql, {'collection_id': collection_id}) return [an[0] for an in attribute_names] if attribute_names else [] def get_attribute_id(schema, attribute_name: str): """ Gets collection type """ sql = f"""SELECT a_id FROM commons.system_label WHERE a_system_name = %(attribute_name)s;""" for attribute_id in run_query(sql, {'attribute_name': attribute_name}): return attribute_id[0] return None def get_collection_w_profiles(schema, collection_id): for collection in run_query(f"""SELECT id, name, status, parent_id AS "parentId" , TO_CHAR(date_created, 'YYYY-MM-DD"T"HH24:MI:SS.US"Z"') AS "dateCreated" , collection_type AS "type" , (SELECT count(*) FROM {schema}.collection_fan CF WHERE CF.collection_id=C.id) AS "totalProfiles" FROM {schema}.collection C WHERE C.id = %(collection_id)s""", {'collection_id': collection_id}, return_type="dict"): return collection else: raise RuntimeError('Invalid collection') def change_collection_status(schema: str, user_id: str, collection_id: int, status: str, collection_type: str = None, collection_backup=None): return update_collection(schema, user_id, collection_id, status=status, collection_type=collection_type, collection_backup=collection_backup) def update_collection(schema: str, user_id: str, collection_id: int, name: str = None, collection_type: str = None, status: str = None, collection_backup=None): """ Updating the 'status' field of collection """ sql_params: Dict[str, Union[str, int]] = {} if name: sql_params['name'] = name if status: sql_params['status'] = status if collection_type: sql_params['collection_type'] = collection_type if sql_params: sql = f"UPDATE {schema}.collection SET {','.join([f'{q} = %({q})s' for q in sql_params])} " \ f"WHERE id = %(collection_id)s RETURNING source, collection_type, " \ f"(SELECT status FROM {schema}.collection WHERE id = %(collection_id)s)" sql_params['collection_id'] = int(collection_id) update_result = run_query(sql, sql_params) # can be already deleted, so returns nothing try: try: collection = get_collection_w_profiles(schema, collection_id) except RuntimeError: # can be already deleted, so returns nothing collection = { 'id': collection_id, 'status': status, 'parentId': collection_backup['parentId'] if collection_backup and 'parentId' in collection_backup else None, 'name': collection_backup['name'] if collection_backup and 'name' in collection_backup else None, 'dateCreated': collection_backup['dateCreated'] if collection_backup and 'dateCreated' in collection_backup else None } if not status or not update_result or status != update_result[0][2]: # Only notify and log, when status actually changed ct = update_result[0][1] if update_result else \ collection_backup['type'] if collection_backup and 'type' in collection_backup else \ None parent_id = collection['parentId'] if 'parentId' in collection else None if ct in VISIBLE_COLLECTIONS and \ (parent_id is None or status in ('finished', 'deleted')): # Only care about visible collection types or finished status for children. try: # Send collection fields defined in appsync-serverless to AppSync if schema.startswith('c') or schema.startswith('w'): # workspace schema appsync.send_collection(schema, collection) elif schema.startswith('a'): # alliance schema appsync.send_alliance(schema, {'id': schema, 'status': status}) except Exception as e: component_logger.exception(e) # As a last step, store status log if status: sql_params['collection_id'] = int(collection_id) sql_params['user_id'] = user_id sql = f"INSERT INTO {schema}.collection_status_log (collection_id, user_id, status) " \ f"VALUES (%(collection_id)s, %(user_id)s, %(status)s ) " run_query(sql, sql_params) add_event_log(schema, collection_id, user_id, 'collection_delete' if status == 'deleted' else 'collection_status_change', status) if name: add_event_log(schema, collection_id, user_id, 'collection_name_change', name) return collection except Exception as e: component_logger.exception(e) return {} TYPE_MAP = {'int64': 'bigint', 'object': 'varchar(255)', 'float64': 'float'} def string_to_db(df: pd.DataFrame, schema, table, specify_columns=False, create_table=False) -> None: """ This method can be used to loading bigger dataframes in order to increase loading speed df has to have exact structure as the target DB table NOTE: can't use serial and other default db behaviour """ with get_rds_connection() as connection: with connection.cursor() as cursor: f = StringIO() df.to_csv(f, index=False, header=False, sep='\t') # remove header f.seek(0) # move position to beginning of file before reading if create_table: cursor.execute(f"CREATE TABLE {schema}.{table} ({','.join(f'{col} {TYPE_MAP[str(df.dtypes[col])]}' for col in df.columns)});") if specify_columns: column_spec = f"({','.join(df.columns)})" else: column_spec = '' cursor.copy_expert(f"COPY {schema}.{table} {column_spec} FROM STDIN (FORMAT csv, DELIMITER '\t', HEADER False)", f) def write_csv_to_efs(df, schema, table, header=False, attributes=None) -> str: """ Writes out data frame to a file """ if attributes is None: attributes = [] out_path = f"{EFS_PATH}/{schema}-{table}-{'-'.join(attributes)}-{make_iso_date_now()}.csv" with open(f"{out_path}", 'x', encoding='utf-8') as out: df.to_csv(out, index=False, header=header, sep='\t') return out_path def get_related_collections(schema, collection_id: Union[int, List[int]], only_types=None, include_segment_sources=False, no_parents=False) -> list: """Returns a list of all collection_ids related with collection_id including the input collection_id include_segment_sources=True causes recursive collection of segment sources until real datasource collections are found """ only_types = only_types or [] internal_only_types = only_types + ['segment'] if include_segment_sources and \ only_types and 'segment' not in only_types else only_types if isinstance(collection_id, (int, str)): collection_ids = [int(collection_id)] elif isinstance(collection_id, list): collection_ids = [int(cid) for cid in collection_id] else: raise RuntimeError('Check your inputs, please') if len(collection_ids) == 0 or (len(collection_ids) == 1 and collection_ids[0] == -1) or -1 in collection_ids: parent_part = f"parent as (SELECT id FROM {schema}.collection " \ f"WHERE collection_type = 'source')" else: if no_parents: parent_part = f"parent as (SELECT DISTINCT id " \ f"FROM {schema}.collection " \ f"WHERE id IN ({','.join(str(cid) for cid in collection_ids)}))" elif include_segment_sources: parent_part = f"RECURSIVE pre_parent as (SELECT DISTINCT id " \ f"FROM {schema}.collection " \ f"WHERE id IN ({','.join(str(cid) for cid in collection_ids)}) " \ f"UNION SELECT c.parent_id FROM {schema}.collection c INNER JOIN pre_parent p ON p.id = c.id AND c.parent_id is not NULL), " \ f"parent as (SELECT DISTINCT id " \ f"FROM {schema}.collection " \ f"WHERE id IN (SELECT * from pre_parent) " \ f"UNION SELECT sc.collection_id as id FROM {schema}.set_collection sc INNER JOIN parent p ON p.id = sc.set_id)" else: parent_part = f"RECURSIVE parent as (SELECT DISTINCT id " \ f"FROM {schema}.collection " \ f"WHERE id IN ({','.join(str(cid) for cid in collection_ids)}) " \ f"UNION SELECT c.parent_id FROM {schema}.collection c INNER JOIN parent p ON p.id = c.id AND c.parent_id is not NULL)" data_types = f""" AND collection_type IN ({','.join([f"'{tt}'" for tt in internal_only_types])})""" if internal_only_types else '' query = f'''WITH {parent_part} SELECT id, collection_type, parent_id FROM {schema}.collection WHERE (id IN (SELECT * FROM parent) OR parent_id IN (SELECT * FROM parent)){data_types};''' result = pd_read_sql(query) try: segments = result.loc[(result["collection_type"] == "segment")] if len(segments): # remove all algos that are not directly under the earliest set, if there are some later ones. min_segment = segments['id'].min() if not result[(result['parent_id'] == min_segment) & (result["collection_type"] == "algo")].empty: result.drop(result[(result['parent_id'] != min_segment) & (result["collection_type"] == "algo")].index, inplace=True) except Exception as e: component_logger.exception(e) try: if only_types: return list(result.loc[result['collection_type'].isin(only_types)]['id']) else: return list(result['id']) except: return [] def get_parent(schema, collection_id): query = f'SELECT COALESCE(parent_id, {collection_id}) FROM {schema}.collection WHERE id = {collection_id}' res = run_query(query, {'collection_id': collection_id}) if res: return res[0][0] else: return None def get_table_structure(schema, table_name) -> list: """ Returns a list of column names in specific table""" sql = f""" SELECT column_name FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_NAME = '{table_name}' AND table_schema = '{schema}'; """ table_structure = pd_read_sql(sql)['column_name'] return table_structure def refresh_materialized_queries(schema): rds_query([f"REFRESH MATERIALIZED VIEW {schema}.fan_collection_attribute;", f"REFRESH MATERIALIZED VIEW {schema}.fan_attribute_count;"]) return True