-
Notifications
You must be signed in to change notification settings - Fork 20
/
prerun.py
220 lines (170 loc) · 6.18 KB
/
prerun.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
import os
import sys
import subprocess
import psycopg2
try:
from urllib.request import urlopen
from urllib.error import URLError
except ImportError:
from urllib2 import urlopen
from urllib2 import URLError
import time
import re
import json
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
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
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}/schema/name?wt=json'.format(url=url)
try:
connection = urlopen(search_url)
except URLError as e:
print(str(e))
print("[prerun] Unable to connect to solr, waiting...")
time.sleep(10)
check_solr_connection(retry=retry - 1)
else:
import re
conn_info = connection.read()
schema_name = json.loads(conn_info)
if 'ckan' in schema_name['name']:
print('[prerun] Succesfully connected to solr and CKAN schema loaded')
else:
print('[prerun] Succesfully connected to solr, but CKAN schema not found')
def init_db():
db_command = ["ckan", "-c", ckan_ini, "db", "upgrade"]
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()
# Remove internal pg command as psycopg2 does not like it
perms_sql = re.sub(b'\\\\connect "(.*)"', b"", 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())
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 b"User:None" not in re.sub(b"\s", b"", out):
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))
# cleanup permissions
# We're running as root before pivoting to uwsgi and dropping privs
data_dir = "%s/storage" % os.environ['CKAN_STORAGE_PATH']
command = ["chown", "-R", "ckan:ckan", data_dir]
subprocess.call(command)
print("[prerun] Ensured storage directory is owned by ckan")
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()