raena-crawler-engine/tokopedia_crawler_engine/tokopedia_db_migrations.py

from tokopedia_logger import logger
from tokopedia_db_writer import DB
from tokopedia_config import Config

class db_migrations():
  config = Config().get()

  def __init__(self):
    logger.info('Running database migrations')
    self.updateSource()
    self.updateCategoryColumn()
    self.alterCrawlerTracker()
    self.alterProductTab()
    self.alterResellerStoreTab()
    logger.info('Database migrations completed')

  def updateSource(self):
    table = f'{self.config.get("crawler_schema")}.{self.config.get("source_tab")}'
    target = self.config.get("crawler_target")
    target_url = self.config.get("crawler_target_url")
    query = f'''INSERT INTO {table} (source_name, source_main_url)
      SELECT '{target}', '{target_url}'
      WHERE
          NOT EXISTS (
              SELECT id FROM {table} WHERE source_name = '{target}'
          );'''
    try:
      DB().execute_query(query)
    except:
      logger.error(f'Problem while creating source in {table}')
      exit(1)

  def updateCategoryColumn(self):
    table = f'{self.config.get("crawler_schema")}.{self.config.get("category_tab")}'
    aud_table = f'{self.config.get("crawler_schema")}.aud_{self.config.get("category_tab")}'
    query = f'Alter table {table} ADD COLUMN IF NOT EXISTS category_slug character varying UNIQUE'
    aud_query = f'Alter table {aud_table} ADD COLUMN IF NOT EXISTS category_slug character varying UNIQUE'
    try:
      DB().execute_query(query)
      DB().execute_query(aud_query)
    except:
      logger.error(f'Problem while updating column in {table}')
      exit(1)

  def alterCrawlerTracker(self):
    table = f'{self.config.get("crawler_schema")}.{self.config.get("tracker_tab")}'
    query = f'''
      ALTER TABLE {table}
      ADD CONSTRAINT unique_product_page_url UNIQUE (product_page_url);
    '''
    try:
      DB().execute_query(query)
    except:
      # This might be the reason of a silent error
      pass

  def alterProductTab(self):
    table = f'{self.config.get("crawler_schema")}.{self.config.get("product_tab")}'
    aud_table = f'{self.config.get("crawler_schema")}.aud_{self.config.get("product_tab")}'
    query = f'Alter table {table} ADD COLUMN IF NOT EXISTS rce_source_id bigint;'
    aud_query = f'Alter table {aud_table} ADD COLUMN IF NOT EXISTS rce_source_id bigint;'

    constraint_query = f'''
      ALTER TABLE {table}
      ADD CONSTRAINT product_source_id_ukey UNIQUE (rce_source_product_id, rce_source_id);
    '''

    try:
      DB().execute_query(query + aud_query)
    except:
      logger.error(f'Problem while updating column in {table}')
      exit(1)

    try:
      DB().execute_query(constraint_query)
    except:
      pass

  def alterResellerStoreTab(self):
    table = f'{self.config.get("crawler_schema")}.{self.config.get("reseller_store_tab")}'
    aud_table = f'{self.config.get("crawler_schema")}.aud_{self.config.get("reseller_store_tab")}'
    query = f'Alter table {table} ADD COLUMN IF NOT EXISTS rce_source_id bigint;'
    aud_query = f'Alter table {aud_table} ADD COLUMN IF NOT EXISTS rce_source_id bigint;'

    constraint_query = f'''
      ALTER TABLE {table}
      ADD CONSTRAINT store_source_id_ukey UNIQUE (rce_source_store_id, rce_source_id);
    '''

    aud_constraint_query = f'''
      ALTER TABLE {aud_table}
      ADD CONSTRAINT aud_store_source_id_ukey UNIQUE (rce_source_store_id, rce_source_id);
    '''

    try:
      DB().execute_query(query + aud_query)
    except:
      logger.error(f'Problem while updating column in {table}')
      exit(1)

    try:
      DB().execute_query(constraint_query)
    except:
      pass

    try:
      DB().execute_query(aud_constraint_query)
    except:
      pass