import os import sys import subprocess import psycopg2 import urllib3 import time import re ckan_ini = os.environ.get('CKAN_INI', '/srv/app/ckan.ini') RETRY = 5 def update_plugins(): plugins = os.environ.get('CKAN__PLUGINS', '') print('[prerun] Setting the following plugins in {}:'.format(ckan_ini)) print(plugins) cmd = ['ckan', 'config-tool', ckan_ini, 'ckan.plugins = {}'.format(plugins)] subprocess.check_output(cmd, stderr=subprocess.STDOUT) print('[prerun] Plugins set.') def check_main_db_connection(retry=None): conn_str = os.environ.get('CKAN_SQLALCHEMY_URL') if not conn_str: print('[prerun] CKAN_SQLALCHEMY_URL not defined, not checking db') return check_db_connection(conn_str, retry) def check_datastore_db_connection(retry=None): conn_str = os.environ.get('CKAN_DATASTORE_WRITE_URL') if not conn_str: print('[prerun] CKAN_DATASTORE_WRITE_URL not defined, not checking db') return check_db_connection(conn_str, retry) def check_db_connection(conn_str, retry=None): if retry is None: retry = RETRY elif retry == 0: print('[prerun] Giving up after 5 tries...') sys.exit(1) try: connection = psycopg2.connect(conn_str) except psycopg2.Error as e: print(str(e)) print('[prerun] Unable to connect to the database, waiting...') time.sleep(10) check_db_connection(conn_str, retry=retry - 1) else: connection.close() def check_solr_connection(retry=None): if retry is None: retry = RETRY elif retry == 0: print('[prerun] Giving up after 5 tries...') sys.exit(1) url = os.environ.get('CKAN_SOLR_URL', '') search_url = '{url}/select/?q=*&wt=json'.format(url=url) http = urllib3.PoolManager() try: r = http.request('GET', search_url) except urllib3.exceptions.ConnectionError as e: print(str(e)) print('[prerun] Unable to connect to solr, waiting...') time.sleep(10) check_solr_connection(retry=retry - 1) else: print('[prerun] Connection Status from SOLR is ', (r.status)) def init_db(): db_command = ['ckan', '-c', ckan_ini, 'db', 'init'] print('[prerun] Initializing or upgrading db - start') try: subprocess.check_output(db_command, stderr=subprocess.STDOUT) print('[prerun] Initializing or upgrading db - end') except subprocess.CalledProcessError as e: if 'OperationalError' in e.output: print(e.output) print('[prerun] Database not ready, waiting a bit before exit...') time.sleep(5) sys.exit(1) else: print(e.output) raise e def init_datastore_db(): conn_str = os.environ.get('CKAN_DATASTORE_WRITE_URL') if not conn_str: print('[prerun] Skipping datastore initialization') return datastore_perms_command = ['ckan', '-c', ckan_ini, 'datastore', 'set-permissions'] connection = psycopg2.connect(conn_str) cursor = connection.cursor() print('[prerun] Initializing datastore db - start') try: datastore_perms = subprocess.Popen( datastore_perms_command, stdout=subprocess.PIPE) perms_sql = datastore_perms.stdout.read().decode('utf-8') # Remove internal pg command as psycopg2 does not like it perms_sql = re.sub('\\\\connect \"(.*)\"', '', perms_sql) cursor.execute(perms_sql) for notice in connection.notices: print(notice) connection.commit() print('[prerun] Initializing datastore db - end') print(datastore_perms.stdout.read().decode('utf-8')) except psycopg2.Error as e: print('[prerun] Could not initialize datastore') print(str(e)) except subprocess.CalledProcessError as e: if 'OperationalError' in e.output: print(e.output) print('[prerun] Database not ready, waiting a bit before exit...') time.sleep(5) sys.exit(1) else: print(e.output) raise e finally: cursor.close() connection.close() def create_sysadmin(): name = os.environ.get('CKAN_SYSADMIN_NAME') password = os.environ.get('CKAN_SYSADMIN_PASSWORD') email = os.environ.get('CKAN_SYSADMIN_EMAIL') if name and password and email: # Check if user exists command = ['ckan', '-c', ckan_ini, 'user', 'show', name,] out = subprocess.check_output(command) if 'User:None' not in re.sub(r'\s', '', out.decode()): print('[prerun] Sysadmin user exists, skipping creation') return # Create user command = ['ckan', '-c', ckan_ini, 'user', 'add', name, 'password=' + password, 'email=' + email] subprocess.call(command) print('[prerun] Created user {0}'.format(name)) # Make it sysadmin command = ['ckan', '-c', ckan_ini, 'sysadmin', 'add', name] subprocess.call(command) print('[prerun] Made user {0} a sysadmin'.format(name)) if __name__ == '__main__': maintenance = os.environ.get('MAINTENANCE_MODE', '').lower() == 'true' if maintenance: print('[prerun] Maintenance mode, skipping setup...') else: check_main_db_connection() init_db() update_plugins() check_datastore_db_connection() init_datastore_db() check_solr_connection() create_sysadmin()