import os import json import boto3 import psycopg2 from psycopg2._psycopg import ProgrammingError from psycopg2.extras import RealDictCursor import logging log = logging.getLogger() log.setLevel(logging.INFO) REGION = os.environ["AWS_REGION"] if "AWS_REGION" in os.environ else "eu-west-1" S3_BUCKET_DEVEL = "frontend-api-devel-filestore" def get_secret(secret_name): # Create a Secrets Manager client session = boto3.session.Session() client = session.client( service_name='secretsmanager', region_name=REGION ) get_secret_value_response = json.loads(client.get_secret_value(SecretId=secret_name)['SecretString']) return get_secret_value_response DB_CONN = {'fansifter-rds': None, 'fansifter-rds-test': None, 'fansifter-rds-live': None, 'sandbox': None} DB_PORT = {'fansifter-rds': 5432, 'fansifter-rds-test': 5433, 'fansifter-rds-live': 5434, 'sandbox': None} def get_rds_connection(conn_name): global DB_CONN if DB_CONN[conn_name] is None: try: secret = get_secret(conn_name) # Toomas: this port mapping only exists for ssh tunneling setup in my virtual machine. # Anywhere else you can just use secret['port'] instead of DB_PORT[conn_name] in get_rds_connection() DB_CONN[conn_name] = psycopg2.connect(host=secret['host'], #'localhost', port=DB_PORT[conn_name], database=secret['dbname'], user=secret['username'], password=secret['password'], cursor_factory=RealDictCursor) except Exception as e: log.exception(e) return DB_CONN[conn_name] def rds_query(connection_name, query, vars=None, fetch=False): with get_rds_connection(connection_name) as conn: with conn.cursor() as cur: if isinstance(query, list): for statement in query: cur.execute(statement, vars=vars) else: cur.execute(query, vars=vars) try: if fetch: return cur.fetchall() except ProgrammingError as e: return [] # select here which database connections should be included. databases = [ 'fansifter-rds', # 'fansifter-rds-test', # 'fansifter-rds-live', ] only_schemas = ['c6fe2499ee0e679030659d9bb787ff732ed3da6721d05f96f615a2051'] new_version = 'mvp 2.2' def initialize_schema(connection, schema, user_id, caw, email): # schema_name: str = create_workspace_schema(client_id) fb_base_tables = f""" CREATE TABLE IF NOT EXISTS {schema}.fb_adaccount ( id text, business_id text, business_name text, fb_user_id text, user_id text, name text, description text, status int, timestamp_created TIMESTAMP, timestamp_updated TIMESTAMP ); ALTER TABLE {schema}.fb_adaccount DROP CONSTRAINT IF EXISTS unique_fb_adaccount_id; ALTER TABLE {schema}.fb_adaccount ADD CONSTRAINT unique_fb_adaccount_id UNIQUE(id); CREATE TABLE IF NOT EXISTS {schema}.fb_audience ( id serial, external_id varchar(255), collection_id int, account_id varchar(255), user_id varchar(40), type varchar(255), parent_id varchar(255), status varchar(255), name varchar(255), description varchar(255), approximate_count int, actual_count int, delivery_status_code varchar(255), delivery_status_description varchar(255), customer_file_source varchar(255), operation_status_code varchar(255), operation_status_description varchar(255), ratio json, -- {{'id': '0.01', 'name': '1'}} countries json, -- [{{'id': 'US', 'name': 'United States'}}] timestamp_created TIMESTAMP, timestamp_last_updated TIMESTAMP ); CREATE UNIQUE INDEX IF NOT EXISTS idx_fb_audience ON {schema}.fb_audience(external_id); CREATE TABLE IF NOT EXISTS {schema}.fb_audience_shared_state ( id varchar(255), audience_id int, adaccount_id varchar(255), shared_status varchar(255), timestamp_created TIMESTAMP, timestamp_last_updated TIMESTAMP ); CREATE UNIQUE INDEX IF NOT EXISTS idx_fb_shared_shared_state ON {schema}.fb_audience(id); INSERT INTO {schema}.meta_data (key,val) VALUES ('fb_base_version', 'mvp 2.0') ON CONFLICT (key) DO NOTHING; CREATE TABLE IF NOT EXISTS {schema}.fb_temp_audience_data ( EMAIL varchar(255), DOB varchar(255), GEN varchar(255), FN varchar(255), CT varchar(255), COUNTRY varchar(255), audience_id varchar(255) ); """ rds_query(connection, fb_base_tables) query = f"""INSERT INTO {schema}.meta_data (key,val) VALUES ('creator', %(user_id)s) ON CONFLICT (key) DO NOTHING;""" rds_query(connection, query, {'user_id': user_id}, fetch=False) if caw == 'company': create_management_tables = f""" CREATE TABLE IF NOT EXISTS {schema}.company ( company_id varchar(60), name varchar(255), current_package_id int, subscription_valid_until_date date ); CREATE TABLE IF NOT EXISTS {schema}.product_package_purchase_history ( id serial, start_date timestamp, end_date timestamp, package_id int, payment_processor_token varchar, paid_amount numeric -- how much did client pay for this package ); CREATE TABLE IF NOT EXISTS {schema}.company_user ( user_id varchar(40), name varchar(255), email varchar(255), selected_worspace varchar(60) ); CREATE UNIQUE INDEX IF NOT EXISTS idx_company_user_id ON {schema}.company_user(user_id); CREATE UNIQUE INDEX IF NOT EXISTS idx_company_user_email ON {schema}.company_user(email); CREATE TABLE IF NOT EXISTS {schema}.user_roles ( id varchar(60), company_alliance_workspace varchar(20), user_id varchar(40), roles json -- list of role names ); CREATE UNIQUE INDEX IF NOT EXISTS idx_user_access ON {schema}.user_roles(id, user_id); CREATE UNIQUE INDEX IF NOT EXISTS idx_user_caw ON {schema}.user_roles(user_id, company_alliance_workspace, id); CREATE SEQUENCE IF NOT EXISTS {schema}.alliance_num_seq START WITH 1 INCREMENT BY 1; CREATE TABLE IF NOT EXISTS {schema}.alliance ( id varchar(60), num integer NOT NULL DEFAULT nextval('{schema}.alliance_num_seq') -- creating serial ); ALTER SEQUENCE {schema}.attribute_id_seq OWNED BY {schema}.alliance.num; CREATE UNIQUE INDEX IF NOT EXISTS idx_alliance_id ON {schema}.alliance(id); CREATE UNIQUE INDEX IF NOT EXISTS idx_alliance_num ON {schema}.alliance(num); CREATE SEQUENCE IF NOT EXISTS {schema}.workspace_num_seq START WITH 1 INCREMENT BY 1; CREATE TABLE IF NOT EXISTS {schema}.workspace ( id varchar(60), num integer NOT NULL DEFAULT nextval('{schema}.workspace_num_seq') -- creating serial ); ALTER SEQUENCE {schema}.workspace_num_seq OWNED BY {schema}.workspace.num; CREATE UNIQUE INDEX IF NOT EXISTS idx_workspace_id ON {schema}.workspace(id); CREATE UNIQUE INDEX IF NOT EXISTS idx_workspace_num ON {schema}.workspace(num); INSERT INTO {schema}.workspace (id) VALUES ('{schema}') ON CONFLICT (id) DO NOTHING; INSERT INTO {schema}.meta_data (key,val) VALUES ('management_version', 'mvp 2.0') ON CONFLICT (key) DO NOTHING; """ rds_query(connection, create_management_tables) create_fb_company_tables = f""" CREATE TABLE IF NOT EXISTS {schema}.fb_user ( user_id text CONSTRAINT unique_fb_user_id UNIQUE, access_token text, token_type text, expires_in int, token_expiration_timestamp TIMESTAMP, signed_request text, graph_domain text, granted_permissions text, data_access_expiration_int int, data_access_expiration_timestamp TIMESTAMP, status text, timestamp_created TIMESTAMP, timestamp_last_updated TIMESTAMP ); CREATE TABLE IF NOT EXISTS {schema}.fb_sdk_adaccount ( id text CONSTRAINT unique_fb_sdk_adaccount_id UNIQUE, business_id text, business_name text, fb_user_id text, user_id text, name text, description text, status int, timestamp_created TIMESTAMP, timestamp_updated TIMESTAMP ); INSERT INTO {schema}.meta_data (key,val) VALUES ('fb_company_version', 'mvp 2.0') ON CONFLICT (key) DO NOTHING; """ rds_query(connection, create_fb_company_tables) company_name = '' query = f"""INSERT INTO commons.user_company (user_id,email, default_schema, company_name) VALUES (%(user_id)s, %(email)s, %(schema)s, %(company_name)s) ON CONFLICT (user_id) DO UPDATE SET email = excluded.email, default_schema = excluded.default_schema;""" rds_query(connection, query, {'user_id': user_id, 'schema': schema, 'email': email, 'company_name': company_name}, fetch=False) 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;""" rds_query(connection, query, {'user_id': user_id, 'schema': schema}, fetch=False) # 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 = 'lite'));""" # rds_query(connection, query) return schema other_buckets = set() for connection in databases: # all_schemas = rds_query(connection, # "select table_schema, table_name from information_schema.tables where table_name IN ('meta_data' )", fetch=True) all_schemas = rds_query(connection, "select default_schema, user_id, email from commons.user_company", fetch=True) print(f"Found {len(all_schemas)} schemas in database {connection}") for schema, user_id, email in [sch.values() for sch in all_schemas]: print(f'Working on: {schema} {user_id} {email}') # if schema in only_schemas: # continue try: ws = rds_query(connection, f"SELECT fb_user_id FROM {schema}.fb_sdk_adaccount", fetch=True) ws = rds_query(connection, f"SELECT * FROM {schema}.fb_audience", fetch=True) except Exception as e: print('OOPS', e) try: rds_query(connection, f"DROP TABLE {schema}.fb_sdk_adaccount") except: pass initialize_schema(connection, schema, user_id, 'company', email) continue try: ver = rds_query(connection, f"SELECT val FROM {schema}.meta_data WHERE key = 'version'", fetch=True)[0]['val'] except Exception as e: ver = None if False: statements=[ #f"alter table {schema}.fb_audience alter column parent_id type int using parent_id::int;" f"alter table {schema}.fb_audience alter column parent_id type varchar(255) using parent_id::varchar(255);" ] try: #print(start_num) rds_query(connection, statements) # , vars={'new_version': new_version}) except Exception as e: log.exception("oops", exc_info=e) else: if ver == 'mvp 2.2': print('Already version mvp 2.2') if ver is None: print('Too old to upgrade')