database.init_db

  1"""Initialisation and deinitialisation of the database."""
  2
  3from logging import Logger, getLogger
  4from logging.config import dictConfig
  5from pathlib import Path
  6from re import IGNORECASE, search, split
  7from string.templatelib import Template
  8
  9from internal.settings.env import database_settings
 10from psycopg import Connection, Error, connect
 11from uvicorn.config import LOGGING_CONFIG
 12
 13from database.db_constants import ALLERGENS, ANALYTICS_GRAPHS_TYPES, BADGES, CATEGORIES
 14
 15dictConfig(LOGGING_CONFIG)
 16SCHEMA_PATH = Path("database/migrations/schema.sql")
 17
 18
 19def load_queries() -> list[str]:
 20    """Loads table creation queries from schema.
 21
 22    Returns:
 23      list of separate queries
 24    """
 25    content = SCHEMA_PATH.read_text(encoding="utf-8")
 26    raw_queries = split(r"\n\s*\n", content)
 27    return [q.strip() for q in raw_queries if q.strip()]
 28
 29
 30# Connect to the database
 31def get_db_connection(logger: Logger) -> Connection | None:
 32    """Opens connection to the database.
 33
 34    Args:
 35      logger: configured logger
 36
 37    Returns:
 38      Opened connection
 39    """
 40    try:
 41        connection = connect(
 42            host=database_settings.host,
 43            port=database_settings.port,
 44            dbname=database_settings.database,
 45            user=database_settings.username,
 46            password=database_settings.password,
 47        )
 48        if connection.closed == 0:
 49            logger.info("Connection to database successful.")
 50            return connection
 51        logger.error("Connection is closed")
 52        return None
 53    except Error as e:
 54        logger.error(f"Error while connecting to Postgres: {e}")
 55        return None
 56
 57
 58# Get entity name and type
 59def get_type_and_name(sql: str) -> tuple[str, str] | None:
 60    """Extracts name and type of the entity from the query.
 61
 62    Args:
 63        sql: sql string
 64
 65    Returns:
 66      tuple with table type and name
 67    """
 68    table_match = search(
 69        r"CREATE\s+TABLE\s+(?:IF\s+NOT\s+EXISTS\s+)?`?(\w+)`?", sql, IGNORECASE
 70    )
 71    if table_match:
 72        return ("TABLE", table_match.group(1))
 73
 74    enum_match = search(r"CREATE\s+TYPE\s+`?(\w+)`?\s+AS\s+ENUM", sql, IGNORECASE)
 75    if enum_match:
 76        return ("ENUM", enum_match.group(1))
 77    return None
 78
 79
 80def create_all_tables(
 81    logger: Logger, table_queries: list[str], conn: Connection
 82) -> None:
 83    """Execute sql for each table.
 84
 85    Args:
 86      logger: configured logger
 87      table_queries: list of the queries for table craetion
 88      conn: connection to the database
 89    """
 90    for sql in table_queries:
 91        type_name = get_type_and_name(sql)
 92        if not type_name:
 93            logger.error("unrecognised entity type")
 94            continue
 95        try:
 96            with conn.cursor() as cursor:
 97                cursor.execute(Template(sql))
 98                logger.info(f"{type_name[0].lower()} {type_name[1]} created.")
 99        except Error as err:
100            logger.error(f"Failed to create {type_name[0]} {type_name[1]}: {err}")
101        conn.commit()
102
103
104def show_all_tables(logger: Logger, conn: Connection) -> None:
105    """Show all tables located in the database.
106
107    Args:
108      logger: configured logger
109      conn: connection to the database
110    """
111    with conn.cursor() as cursor:
112        cursor.execute("""
113            SELECT table_name
114            FROM information_schema.tables
115            WHERE table_schema NOT IN ('pg_catalog', 'information_schema');
116        """)
117        tables = cursor.fetchall()
118        logger.info("Current tables in the database:")
119        for table in tables:
120            logger.info(table)
121
122
123def drop_all_tables(logger: Logger, table_queries: list[str], conn: Connection) -> None:
124    """Drops all created tables from database.
125
126    Args:
127      logger: configured logger
128      table_queries: list of the queries for table craetion
129      conn: connection to the database
130    """
131    table_queries.reverse()
132    for sql in table_queries:
133        type_name = get_type_and_name(sql)
134        if not type_name:
135            logger.error("unrecognised entity type")
136            continue
137        if type_name[0] == "ENUM":
138            logger.warning(f"skipping {type_name[1]} as enum")
139            continue
140        with conn.cursor() as cursor:
141            cursor.execute(Template(f"DROP {type_name[0]} IF EXISTS {type_name[1]};"))
142            logger.info(f"table {type_name[1]} removed.")
143    conn.commit()
144
145
146def seed_static_data(logger: Logger, conn: Connection) -> None:
147    """Inserts categories and allergens into the db."""
148    try:
149        # Seed Categories
150        logger.info("inserting categories")
151        for category in CATEGORIES:
152            conn.execute(
153                "INSERT INTO category "
154                "(category_id, category_name, category_coefficient)"
155                "VALUES (%s, %s, %s) ON CONFLICT DO NOTHING",
156                (category["cat_id"], category["name"], category["coefficient"]),
157            )
158
159        # Seed Allergens
160        logger.info("inserting allergens")
161        for all_id, name in ALLERGENS.items():
162            conn.execute(
163                "INSERT INTO allergens (allergen_id, allergen_name)"
164                "VALUES (%s, %s) ON CONFLICT DO NOTHING",
165                (all_id, name),
166            )
167        # Seed Bundles
168        logger.info("inserting badges")
169        for badge in BADGES:
170            conn.execute(
171                "INSERT INTO badges (badge_id, name, description)"
172                "VALUES (%s, %s, %s) ON CONFLICT DO NOTHING",
173                (badge["badge_id"], badge["name"], badge["description"]),
174            )
175        logger.info("inserting analytics graphs types")
176        for graphs_types in ANALYTICS_GRAPHS_TYPES:
177            conn.execute(
178                "INSERT INTO analytics_graphs_types "
179                "(graph_type_id, chart_type, graph_summary, x_axis_label, y_axis_label)"
180                "VALUES (%s, %s, %s, %s, %s) ON CONFLICT DO NOTHING",
181                (
182                    graphs_types["graph_type_id"],
183                    graphs_types["chart_type"],
184                    graphs_types["graph_summary"],
185                    graphs_types["x_axis_label"],
186                    graphs_types["y_axis_label"],
187                ),
188            )
189        conn.commit()
190        logger.info("finished inserting data")
191    except Error as e:
192        logger.error(f"Error seeding data: {e}")
193        conn.rollback()
194
195
196def main() -> None:
197    """Entrypoint for database management script."""
198    dictConfig(LOGGING_CONFIG)
199    logger = getLogger("uvicorn.info")
200    table_queries = load_queries()
201    conn = get_db_connection(logger)
202    if not conn:
203        return
204    option = 0
205    while option not in {1, 2, 3, 4}:
206        message = (
207            "(1) Init all table, (2) Show all tables, (3) Drop all tables, (4) exit: "
208        )
209        option = int(input(message))
210    if input(f"confirm option {option} (y/N): ") not in {"y", "Y", "yes"}:
211        option = 0
212    match option:
213        case 1:
214            # Create all tables
215            create_all_tables(logger, table_queries, conn)
216            seed_static_data(logger, conn)
217        case 2:
218            # Show all tables
219            show_all_tables(logger, conn)
220        case 3:
221            # Drop all tables
222            drop_all_tables(logger, table_queries, conn)
223        case 4:
224            logger.info("shutting down: database connection closed")
225    conn.close()
226    logger.info("database connection closed.")
227
228
229if __name__ == "__main__":
230    main()