|
| 1 | +""" Initializes tables in database and adds some sample data for testing |
| 2 | +
|
| 3 | +FIXME: this place does not fill right... see how this script is called |
| 4 | +FIXME: check https://github.com/aio-libs/aiohttp_admin/blob/master/demos/blog/aiohttpdemo_blog/generate_data.py |
| 5 | +FIXME: rename as server.dev.generate_data.py and set dev as an optional sub-package as server[dev] |
| 6 | +
|
| 7 | +
|
| 8 | +Example of usage |
| 9 | +
|
| 10 | + cd services/web/server/tests/mock |
| 11 | + docker-compose up |
| 12 | +
|
| 13 | + cd ../../config |
| 14 | + python init_db.py |
| 15 | +
|
| 16 | +References: |
| 17 | +[1]:https://github.com/aio-libs/aiohttp-demos/blob/master/docs/preparations.rst#environment |
| 18 | +""" |
| 19 | +import logging |
| 20 | +import pathlib |
| 21 | +import sys |
| 22 | + |
| 23 | +from passlib.hash import sha256_crypt |
| 24 | +from sqlalchemy import ( |
| 25 | + MetaData, |
| 26 | + create_engine |
| 27 | +) |
| 28 | +from tenacity import ( |
| 29 | + retry, |
| 30 | + stop_after_attempt, |
| 31 | + wait_fixed |
| 32 | +) |
| 33 | + |
| 34 | +from server.db.model import ( |
| 35 | + permissions, |
| 36 | + users |
| 37 | +) |
| 38 | + |
| 39 | +from server.settings import ( |
| 40 | + config_from_file |
| 41 | +) |
| 42 | + |
| 43 | + |
| 44 | +CURRENT_DIR = pathlib.Path(sys.argv[0] if __name__ == "__main__" else __file__).parent.parent |
| 45 | +CONFIG_DIR = CURRENT_DIR.parent / "config" |
| 46 | + |
| 47 | +logging.basicConfig(level=logging.DEBUG) |
| 48 | +_LOGGER = logging.getLogger(__name__) |
| 49 | + |
| 50 | + |
| 51 | +DSN = "postgresql://{user}:{password}@{host}:{port}/{database}" |
| 52 | + |
| 53 | + |
| 54 | +USER_CONFIG_PATH = CONFIG_DIR / "server.yaml" |
| 55 | +USER_CONFIG = config_from_file(USER_CONFIG_PATH.as_posix()) |
| 56 | +USER_DB_URL = DSN.format(**USER_CONFIG["postgres"]) |
| 57 | +user_engine = create_engine(USER_DB_URL) |
| 58 | + |
| 59 | +TEST_CONFIG_PATH = CONFIG_DIR / "server-test.yaml" |
| 60 | +TEST_CONFIG = config_from_file(TEST_CONFIG_PATH.as_posix()) |
| 61 | +TEST_DB_URL = DSN.format(**TEST_CONFIG["postgres"]) |
| 62 | +test_engine = create_engine(TEST_DB_URL) |
| 63 | + |
| 64 | +# FIXME: admin user/passwords and in sync with other host/port configs |
| 65 | +ADMIN_DB_URL = DSN.format( |
| 66 | + user="postgres", |
| 67 | + password="postgres", |
| 68 | + database="postgres", |
| 69 | + host=USER_CONFIG["postgres"]["host"], |
| 70 | + port=5432 |
| 71 | +) |
| 72 | + |
| 73 | +# TODO: what is isolation_level? |
| 74 | +admin_engine = create_engine(ADMIN_DB_URL, isolation_level="AUTOCOMMIT") |
| 75 | + |
| 76 | + |
| 77 | +def setup_db(config): |
| 78 | + db_name = config["database"] |
| 79 | + db_user = config["user"] |
| 80 | + db_pass = config["password"] |
| 81 | + |
| 82 | + # TODO: compose using query semantics. Clarify pros/cons vs string cli? |
| 83 | + conn = admin_engine.connect() |
| 84 | + conn.execute("DROP DATABASE IF EXISTS %s" % db_name) |
| 85 | + conn.execute("DROP ROLE IF EXISTS %s" % db_user) |
| 86 | + conn.execute("CREATE USER %s WITH PASSWORD '%s'" % (db_user, db_pass)) |
| 87 | + conn.execute("CREATE DATABASE %s ENCODING 'UTF8'" % db_name) |
| 88 | + conn.execute("GRANT ALL PRIVILEGES ON DATABASE %s TO %s" % |
| 89 | + (db_name, db_user)) |
| 90 | + conn.close() |
| 91 | + |
| 92 | + |
| 93 | +def teardown_db(config): |
| 94 | + db_name = config["database"] |
| 95 | + db_user = config["user"] |
| 96 | + |
| 97 | + conn = admin_engine.connect() |
| 98 | + conn.execute(""" |
| 99 | + SELECT pg_terminate_backend(pg_stat_activity.pid) |
| 100 | + FROM pg_stat_activity |
| 101 | + WHERE pg_stat_activity.datname = '%s' |
| 102 | + AND pid <> pg_backend_pid();""" % db_name) |
| 103 | + conn.execute("DROP DATABASE IF EXISTS %s" % db_name) |
| 104 | + conn.execute("DROP ROLE IF EXISTS %s" % db_user) |
| 105 | + conn.close() |
| 106 | + |
| 107 | + |
| 108 | +def create_tables(engine=test_engine): |
| 109 | + meta = MetaData() |
| 110 | + meta.create_all(bind=engine, tables=[users, permissions]) |
| 111 | + |
| 112 | + |
| 113 | +def drop_tables(engine=test_engine): |
| 114 | + meta = MetaData() |
| 115 | + meta.drop_all(bind=engine, tables=[users, permissions]) |
| 116 | + |
| 117 | + |
| 118 | +def sample_data(engine=test_engine): |
| 119 | + generate_password_hash = sha256_crypt.hash |
| 120 | + |
| 121 | + #TODO: use fake to populate database |
| 122 | + # pylint:disable=E1120 |
| 123 | + conn = engine.connect() |
| 124 | + conn.execute(users.insert(), [ |
| 125 | + |
| 126 | + "passwd": generate_password_hash("z43"), |
| 127 | + "is_superuser": False, |
| 128 | + "disabled": False}, |
| 129 | + |
| 130 | + "passwd": generate_password_hash("123"), |
| 131 | + "is_superuser": True, |
| 132 | + "disabled": False}, |
| 133 | + |
| 134 | + "passwd": generate_password_hash("345"), |
| 135 | + "is_superuser": True, |
| 136 | + "disabled": True} |
| 137 | + ]) |
| 138 | + |
| 139 | + conn.execute(permissions.insert(), [ |
| 140 | + {"user_id": 1, |
| 141 | + "perm_name": "tester"}, |
| 142 | + {"user_id": 2, |
| 143 | + "perm_name": "admin"} |
| 144 | + ]) |
| 145 | + |
| 146 | + conn.close() |
| 147 | + |
| 148 | + |
| 149 | +@retry(stop=stop_after_attempt(5), wait=wait_fixed(2)) |
| 150 | +def main(): |
| 151 | + |
| 152 | + config = USER_CONFIG["postgres"] |
| 153 | + engine = user_engine |
| 154 | + |
| 155 | + _LOGGER.info("Setting up db ...") |
| 156 | + setup_db(config) |
| 157 | + _LOGGER.info("") |
| 158 | + |
| 159 | + _LOGGER.info("Creating tables ...") |
| 160 | + create_tables(engine=engine) |
| 161 | + |
| 162 | + _LOGGER.info("Adding sample data ...") |
| 163 | + sample_data(engine=engine) |
| 164 | + |
| 165 | + drop_tables() |
| 166 | + teardown_db(config) |
| 167 | + |
| 168 | + |
| 169 | +if __name__ == "__main__": |
| 170 | + main() |
0 commit comments