Add basic migrations

datalite/migrations.py

@@ -0,0 +1,160 @@
+"""
+Migrations module deals with migrating data when the object
+definitions change. This functions deal with Schema Migrations.
+"""
+from dataclasses import Field
+from os.path import exists
+from typing import Dict, Tuple, List
+import sqlite3 as sql
+
+from .commons import _create_table, _get_table_cols
+
+
+def _get_db_table(class_: type) -> Tuple[str, str]:
+    """
+    Check if the class is a datalite class, the database exists
+    and the table exists. Return database and table names.
+
+    :param class_: A datalite class.
+    :return: A tuple of database and table names.
+    """
+    database_name: str = getattr(class_, 'db_path', None)
+    if not database_name:
+        raise TypeError(f"{class_.__name__} is not a datalite class.")
+    table_name: str = class_.__name__.lower()
+    if not exists(database_name):
+        raise FileNotFoundError(f"{database_name} does not exist")
+    with sql.connect(database_name) as con:
+        cur: sql.Cursor = con.cursor()
+        cur.execute(f"SELECT count(*) FROM sqlite_master "
+                    f"WHERE type='table' AND name='{table_name}';")
+        count: int = int(cur.fetchone()[0])
+    if not count:
+        raise FileExistsError(f"Table, {table_name}, already exists.")
+    return database_name, table_name
+
+
+def _get_table_column_names(database_name: str, table_name: str) -> Tuple[str]:
+    """
+    Get the column names of table.
+
+    :param database_name: Name of the database the table
+        resides in.
+    :param table_name: Name of the table.
+    :return: A tuple holding the column names of the table.
+    """
+    with sql.connect(database_name) as con:
+        cur: sql.Cursor = con.cursor()
+        cols: List[str] = _get_table_cols(cur, table_name)
+    return tuple(cols)
+
+
+def _copy_records(database_name: str, table_name: str):
+    """
+    Copy all records from a table.
+
+    :param database_name: Name of the database.
+    :param table_name: Name of the table.
+    :return: A generator holding dataclass asdict representations.
+    """
+    with sql.connect(database_name) as con:
+        cur: sql.Cursor = con.cursor()
+        cur.execute(f'SELECT * FROM {table_name};')
+        values = cur.fetchall()
+        keys = _get_table_cols(cur, table_name)
+        keys.insert(0, 'obj_id')
+    records = (dict(zip(keys, value)) for value in values)
+    return records
+
+
+def _drop_table(database_name: str, table_name: str) -> None:
+    """
+    Drop a table.
+
+    :param database_name: Name of the database.
+    :param table_name: Name of the table to be dropped.
+    :return: None.
+    """
+    with sql.connect(database_name) as con:
+        cur: sql.Cursor = con.cursor()
+        cur.execute(f'DROP TABLE {table_name};')
+        con.commit()
+
+
+def _modify_records(data, col_to_del: Tuple[str], col_to_add: Tuple[str],
+                    flow: Dict[str, str]) -> Tuple[Dict[str, str]]:
+    """
+    Modify the asdict records in accordance
+        with schema migration rules provided.
+
+    :param data: Data kept as asdict in tuple.
+    :param col_to_del: Column names to delete.
+    :param col_to_add: Column names to add.
+    :param flow: A dictionary that explain
+        if the data from a deleted column
+        will be transferred to a column
+        to be added.
+    :return: The modified data records.
+    """
+    records = []
+    for record in data:
+        record_mod = {}
+        for key in record.keys():
+            if key in col_to_del and key in flow:
+                record_mod[flow[key]] = record[key]
+            else:
+                record_mod[key] = record[key]
+        for key_to_add in col_to_add:
+            if key_to_add not in record_mod:
+                record_mod[key_to_add] = None
+        records.append(record_mod)
+    return records
+
+
+def _migrate_records(class_: type, database_name: str, data,
+                     col_to_del: Tuple[str], col_to_add: Tuple[str], flow: Dict[str, str]) -> None:
+    """
+    Migrate the records into the modified table.
+
+    :param class_: Class of the entries.
+    :param database_name: Name of the database.
+    :param data: Data, asdict tuple.
+    :param col_to_del: Columns to be deleted.
+    :param col_to_add: Columns to be added.
+    :param flow: Flow dictionary stating where
+        column data will be transferred.
+    :return: None.
+    """
+    with sql.connect(database_name) as con:
+        cur: sql.Cursor = con.cursor()
+        _create_table(class_, cur, getattr(class_, 'types_table'))
+        con.commit()
+    new_records = _modify_records(data, col_to_del, col_to_add, flow)
+    for record in new_records:
+        del record['obj_id']
+        class_(**record).create_entry()
+
+
+def basic_migrate(class_: type, column_transfer: dict = None) -> None:
+    """
+    Given a class, compare its previous table,
+    delete the fields that no longer exist,
+    create new columns for new fields. If the
+    column_flow parameter is given, migrate elements
+    from previous column to the new ones.
+
+    :param class_: Datalite class to migrate.
+    :param column_transfer: A dictionary showing which
+        columns will be copied to new ones.
+    :return: None.
+    """
+    database_name, table_name = _get_db_table(class_)
+    table_column_names: Tuple[str] = _get_table_column_names(database_name, table_name)
+    values = class_.__dataclass_fields__.values()
+    data_fields: Tuple[Field] = tuple(field for field in values)
+    data_field_names: Tuple[str] = tuple(field.name for field in data_fields)
+    columns_to_be_deleted: Tuple[str] = tuple(column for column in table_column_names if column not in data_field_names)
+    columns_to_be_added: Tuple[str] = tuple(column for column in data_field_names if column not in table_column_names)
+    records = _copy_records(database_name, table_name)
+    _drop_table(database_name, table_name)
+    _migrate_records(class_, database_name, records, columns_to_be_deleted, columns_to_be_added, column_transfer)