import logging from ast import literal_eval from datetime import datetime, timezone from typing import List import pandas as pd from service.conf import settings from service.db import SchemaPath from service.tasks.management.cognito import get_cognito_user_attributes from service.utils import dbmate from service.utils.aws_connectors import df_to_db, run_query logger = logging.getLogger(__name__) def create_workspace_schema(schema_name: str) -> None: dbmate.up( schema_name, schema_path=SchemaPath.WORKSPACE, disable_ssl=settings.MIGRATIONS_DISABLE_SSL, ) def create_management_tables(schema_name: str) -> None: dbmate.up( schema_name, schema_path=SchemaPath.COMPANY, disable_ssl=settings.MIGRATIONS_DISABLE_SSL, ) def initialize_schema( schema: str, user_id: str, caw: str, email: str, is_test=False ) -> str: create_workspace_schema(schema) run_query( f"""INSERT INTO {schema}.meta_data (key,val) VALUES ('creator', %(user_id)s) ON CONFLICT (key) DO NOTHING;""", {"user_id": user_id}, fetch=False, ) if caw == "company": create_management_tables(schema) # This is so automated tests would work. Otherwise, would need to figure # out a way to provide credentials for automated tests. if settings.PROFILE == "auto_test" or is_test: company_name = "test_company" user_name = "test_user" else: pool_name = ( f"frontend-api-{settings.PROFILE}-user-pool" if settings.PROFILE in ["live", "test"] else "frontend-api-devel-user-pool" ) user_attributes = get_cognito_user_attributes(user_id, pool_name) company_name = user_attributes.get("custom:company", "") user_name = " ".join( [ user_attributes.get("given_name", ""), user_attributes.get("family_name", ""), ] ) run_query( """ INSERT INTO commons.user_company (user_id,email, default_schema, company_name, user_name) VALUES (%(user_id)s, %(email)s, %(schema)s, %(company_name)s, %(user_name)s) ON CONFLICT (user_id) DO UPDATE SET email = excluded.email, default_schema = excluded.default_schema;""", { "user_id": user_id, "schema": schema, "email": email, "company_name": company_name, "user_name": user_name, }, fetch=False, ) run_query( f"""INSERT INTO {schema}.user_roles (id, company_alliance_workspace, user_id, roles) VALUES (%(schema)s, 'workspace', %(user_id)s, '["owner"]') ON CONFLICT (id, user_id) DO NOTHING;""", {"user_id": user_id, "schema": schema}, fetch=False, ) run_query( f"""INSERT INTO {schema}.company (company_id, name, current_package_id) VALUES ( '{schema}', (SELECT company_name FROM commons.user_company WHERE default_schema = '{schema}'), (SELECT id FROM commons.packages WHERE package_name = 'default') );""" ) return schema def get_all_schemas_of_path(schema_path: str) -> List[str]: company_schemas = [ schema for [schema] in run_query("SELECT default_schema FROM commons.user_company") ] if schema_path == SchemaPath.COMPANY: return company_schemas elif schema_path == SchemaPath.WORKSPACE: workspace_schemas = [] for company_schema in company_schemas: query = f""" SELECT id FROM {company_schema}.workspace WHERE num > 1; """ company_workspace_schemas = [schema for [schema] in run_query(query)] workspace_schemas.extend(company_workspace_schemas) return workspace_schemas elif schema_path == SchemaPath.ALLIANCE: alliance_schemas = [] for company_schema in company_schemas: query = f""" SELECT id FROM {company_schema}.alliance """ company_alliance_schemas = [schema for [schema] in run_query(query)] alliance_schemas.extend(company_alliance_schemas) # single alliance schema id can be present in multiple workspaces # that share their data to it return list(set(alliance_schemas)) else: raise ValueError(f"Unknown schema type: {schema_path}") def set_company_trial( management_schema, trial_end, package_name, billing_cycle, trial=True, campaign=None ): package_name = package_name.split(" ")[0].capitalize() query = f"""UPDATE {management_schema}.company SET current_package_id = (SELECT id FROM commons.packages WHERE package_name = %(package_name)s AND campaign_name IS NULL), subscription_valid_until_date = %(trial_end)s, unused_campaign_code = %(campaign)s, billing_cycle = %(billing_cycle)s, trial = %(trial)s""" run_query( query, { "package_name": package_name, "trial_end": trial_end, "campaign": campaign, "billing_cycle": billing_cycle, "trial": trial, }, ) def insert_into_upload_url_commons( schema: str, user_id: str, url_path: str, collection_id: int ): current_timestamp = datetime.now(tz=timezone.utc) url_upload_dict = { "schema_name": [schema], "url_path": [url_path], "collection_id": [collection_id], "upload_timestamp": [current_timestamp], "user_id": [user_id], } url_upload_df = pd.DataFrame.from_dict(url_upload_dict) df_to_db(df=url_upload_df, schema="commons", table_name="upload_url") def store_file_labels(source_df, field_map): # TODO! Doesn't store anything ? """ Takes fields from field_map we got from payload, and replaces column names with system fields. Also tries to guess some fields from obvious column names when we don't get them from field_map. Stores the results to collection_labels table too, but this is mostly deprecated already. Sample fieldmap: (old) [ { systemFieldName: 'userEmail', incomingFieldName: 'id' }, { systemFieldName: 'userFirstName', incomingFieldName: 'first_name' }, .... ] Sample fieldmap: (changed 30.06.2020) [ { value: 'userEmail', index: 0 }, { value: 'userFirstName', index: 1 }, .... ] """ old_columns = source_df.columns.to_list() if isinstance(field_map, str): field_map = literal_eval(field_map) """ Map column names with system fields. When we don't have mapping, use existing name """ mapping_dict = {} sys_fields = [] for field in field_map: sys_field = field["value"] # csv_field = old_columns[field['index']] csv_field = old_columns[int(field["index"])] if sys_field: mapping_dict[csv_field] = sys_field sys_fields.append(sys_field) source_df.rename(columns=mapping_dict, inplace=True) new_column_names = list(source_df.columns) # df_to_db(df=collection_labels_df, schema=schema_name, table_name='collection_labels') # logger.info(f"Successfully stored file labels for collection {collection_id}") return source_df, new_column_names, sys_fields def list_caw_schemas( management_schema, workspaces=True, alliances=True, with_caw_type=False ): """Lists all user-accessible caw schemas with with_caw_type == True :returns [(schema, 'workspace'|'alliance')] otherwise just :returns [schema] """ caw_types = [] if workspaces: caw_types.append("workspace") if alliances: caw_types.append("alliance") if not caw_types: raise RuntimeError("Nothing to list") for caw_type in caw_types: for caw in run_query(f"SELECT id FROM {management_schema}.{caw_type}"): if with_caw_type: yield caw[0], caw_type else: yield caw[0]