neuron/migrate.py

#!/bin/sh
#
# NEURON IIoT System for Industry 4.0
# Copyright (C) 2020-2022 EMQ Technologies Co., Ltd All rights reserved.
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation; either
# version 3 of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with this program; if not, write to the Free Software Foundation,
# Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
#

"""":
if type python3 > /dev/null 2>&1
then
    exec python3 "$0" "$@"
elif type python > /dev/null 2>&1
then
    exec python "$0" "$@"
else
    echo >&2 "Python not installed"
fi
exit 1
"""  # """

from __future__ import print_function

import argparse
import os
import re
import sqlite3
import json


VERBOSE = False


def log(message):
    if VERBOSE:
        print(message)


def parse_args():
    """Parse command line arguments."""
    parser = argparse.ArgumentParser()
    parser.add_argument("-v", "--verbose", action="count", help="verbosity")
    parser.add_argument(
        "-s",
        "--schema_dir",
        type=str,
        default="config/",
        help="schema directory",
    )
    parser.add_argument(
        "-d",
        "--data_dir",
        type=str,
        default="persistence/",
        help="data directory",
    )
    subparsers = parser.add_subparsers(
        dest="subcommand", help="supported subcommands"
    )

    import_parser = subparsers.add_parser(
        "import", help="import persistence json data"
    )

    show_parser = subparsers.add_parser(
        "show", help="show database schema version"
    )

    up_parser = subparsers.add_parser(
        "up", help="bring database schema to latest version"
    )
    up_parser.add_argument(
        "--import_once",
        action="store_true",
        help="import persistence json data if database is empty",
    )

    return parser.parse_args()


class Migration:
    def __init__(self, schema_dir="config", database_dir="persistence", database_name="sqlite.db"):
        try:
            database_path = os.path.join(database_dir, database_name)
            conn = sqlite3.connect(database_path)
        except Exception as e:
            print("Connecting database `{}` fail, {}".format(database_path, e))
            raise e

        self.conn_ = conn
        self.cursor_ = self.conn_.cursor()
        self.schema_dir = schema_dir
        self.database_dir_ = database_dir
        self.database_name_ = database_name

        log("Creating migrations table if not exists")
        self.ensure_migrations_table_exists()

    def ensure_migrations_table_exists(self):
        """Create migration table if not exists."""

        self.cursor_.execute(
            """
            CREATE TABLE IF NOT EXISTS migrations (
                migration_id INTEGER PRIMARY KEY,
                version TEXT NOT NULL UNIQUE,
                description TEXT NOT NULL,
                dirty INTEGER NOT NULL,
                created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP)
            """
        )

    def head(self):
        """Read the current version from the database"""

        log("Reading migration head version")
        self.cursor_.execute(
            "SELECT version, description, dirty, created_at FROM migrations"
            " ORDER BY migration_id DESC LIMIT 1"
        )
        row = self.cursor_.fetchone()

        if not row:
            return {
                "version": None,
                "description": None,
                "dirty": None,
                "created_at": None,
            }

        return {
            "version": row[0],
            "description": row[1],
            "dirty": row[2],
            "created_at": row[3],
        }

    def print_head(self):
        head = self.head()

        if not head["version"]:
            print("No version.")
            return

        print(
            "Version: {}, ({} {})".format(
                head["version"], head["created_at"], head["description"]
            )
        )

        if head["dirty"]:
            print("Warning: the migration database is dirty")

    def get_migration_info(self, file_name):
        """Extract the migration info from migration file name."""

        match = re.search(
            "(?P<version>[^_]*)_?(?P<description>.*)\\.sql",
            file_name,
            re.IGNORECASE,
        )

        if not match:
            return {"version": None, "description": None}

        return {
            "version": match.group("version"),
            "description": match.group("description"),
        }

    def should_apply(self, migration_filename, target_version=None):
        """Check whether to apply the migration file."""

        head = self.head()
        current_version = head["version"] or ""

        if target_version and target_version == current_version:
            # already on the target version, don't apply the migration
            return False

        # extract the proposed version from the migration filename
        proposed_version = self.get_migration_info(migration_filename)[
            "version"
        ]

        if not proposed_version:
            # does not have a valid version in the name
            return False

        if proposed_version <= current_version:
            return False
        if target_version and proposed_version > target_version:
            return False

        return True

    def apply(self, import_once):
        """Migrate to the latest state."""
        head = self.head()
        if head["dirty"]:
            print("Latest version is dirty, please tune manually.")
            return

        migration_filenames = list(
            filter(self.should_apply, os.listdir(self.schema_dir))
        )

        # Sort the migrations based on the version encoded in the filename
        migration_filenames.sort()

        for migration_filename in migration_filenames:
            print("Applying migration " + migration_filename)

            info = self.get_migration_info(migration_filename)

            self.cursor_.execute(
                """INSERT INTO migrations (version, description, dirty)
                                   VALUES (?, ?, 1)""",
                (
                    info["version"],
                    info["description"],
                ),
            )
            self.conn_.commit()

            with open(
                os.path.join(self.schema_dir, migration_filename), "r"
            ) as f:
                sql = f.read()
                self.cursor_.executescript(sql)

            self.cursor_.execute(
                "UPDATE migrations SET dirty = 0 WHERE migration_id = ?",
                (self.cursor_.lastrowid,),
            )
            self.conn_.commit()

        if import_once and self.head()["version"] and self.check_no_data():
            self.import_json()

    def check_no_data(self):
        """Check the database contains no data."""
        self.cursor_.execute("SELECT COUNT(*) FROM nodes")
        node_count = self.cursor_.fetchone()[0]
        return node_count == 0

    def import_json(self):
        node_group_names = self.find_json_importable_node_and_group_names()
        node_names = self.find_json_importable_node_names()
        # 1.
        # self.import_plugins_json()
        # 2.
        self.import_nodes_json(node_names)
        # 3.
        self.import_settings_json(node_names)
        # 4.
        self.import_groups_json(node_group_names)
        # 5.
        self.import_tags_json(node_group_names)
        # 6.
        self.import_subscritions_json(node_names)

    def import_plugins_json(self):
        """Import data from persistence JSON files."""
        plugin_file = os.path.join(self.database_dir_, "plugins.json")
        if not os.path.exists(plugin_file):
            return
        data = json.load(open(plugin_file, "r"))
        for library in data["plugins"]:
            log("Import library {}".format(library))
            self.cursor_.execute(
                "INSERT INTO plugins (library) VALUES (?)", (library,)
            )
            self.conn_.commit()

    def import_nodes_json(self, node_names):
        """Import nodes data from persistence JSON files."""
        adapter_file = os.path.join(self.database_dir_, "adapters.json")
        if os.path.exists(adapter_file):
            data = json.load(open(adapter_file, "r"))
            for node in data["nodes"]:
                log("Import node {}".format(node))
                self.cursor_.execute(
                    """INSERT INTO nodes (name, type, state, plugin_name)
                                  VALUES (?, ?, ?, ?)""",
                    (
                        node["name"],
                        1 if node["type"] == 1 else 2,
                        node["state"],
                        node["plugin_name"],
                    ),
                )
            self.conn_.commit()

        for node_name in node_names:
            adapter_file = os.path.join(
                self.database_dir_, "adapters", node_name, "adapter.json"
            )
            if not os.path.exists(adapter_file):
                continue
            data = json.load(open(adapter_file, "r"))
            log("Import node {}".format(data))
            self.cursor_.execute(
                """INSERT INTO nodes (name, type, state, plugin_name)
                              VALUES (?, ?, ?, ?)""",
                (
                    data["name"],
                    data["type"],
                    data["state"],
                    data["plugin_name"],
                ),
            )
            self.conn_.commit()

    def import_settings_json(self, node_names):
        """Import node settings from persistence JSON files."""
        for node_name in node_names:
            settings_file = os.path.join(
                self.database_dir_, "adapters", node_name, "settings.json"
            )

            if not os.path.exists(settings_file):
                continue

            data = open(settings_file, "r").read()
            log("Import setting {}".format(data))
            self.cursor_.execute(
                "INSERT INTO settings (node_name, setting) VALUES(?, ?)",
                (node_name, data),
            )
            self.conn_.commit()

    def import_groups_json(self, node_group_names):
        """Import node groups from persistence JSON files."""
        for node_name, group_name in node_group_names:
            group_file = os.path.join(
                self.database_dir_,
                "adapters",
                node_name,
                "groups",
                group_name,
                "config.json",
            )

            if not os.path.exists(group_file):
                continue

            data = json.load(open(group_file, "r"))
            log("Import group {}".format(data))
            self.cursor_.execute(
                """INSERT INTO groups (driver_name, name, interval)
                               VALUES (?, ?, ?)""",
                (node_name, group_name, data["read_interval"]),
            )
            self.conn_.commit()

    def import_tags_json(self, node_group_names):
        """Import node tags from persistence JSON files."""
        for node_name, group_name in node_group_names:
            tags_file = os.path.join(
                self.database_dir_,
                "adapters",
                node_name,
                "groups",
                group_name,
                "datatags.json",
            )

            if not os.path.exists(tags_file):
                continue

            data = json.load(open(tags_file, "r"))
            for tag in data["tags"]:
                log("Import tag {}".format(tag))
                self.cursor_.execute(
                    """INSERT INTO tags (driver_name, group_name, name, address,
                                         attribute, precision, decimal, type, description)
                                 VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)""",
                    (
                        node_name,
                        group_name,
                        tag["name"],
                        tag["address"],
                        tag["attribute"],
                        0,
                        0,
                        tag["type"],
                        tag.get("description", ""),
                    ),
                )
            self.conn_.commit()

    def import_subscritions_json(self, node_names):
        """Import subscriptions from persistence JSON files."""
        for node_name in node_names:
            subscription_file = os.path.join(
                self.database_dir_, "adapters", node_name, "subscriptions.json"
            )
            if not os.path.exists(subscription_file):
                continue
            data = json.load(open(subscription_file, "r"))
            for subscription in data["subscriptions"]:
                log("Import subscription {}".format(subscription))
                self.cursor_.execute(
                    """INSERT INTO subscriptions (
                            app_name, driver_name, group_name
                       ) VALUES(?, ?, ?)""",
                    (
                        subscription["sub_adapter_name"],
                        subscription["src_adapter_name"],
                        subscription["group_config_name"],
                    ),
                )
            self.conn_.commit()

    def find_json_importable_node_names(self):
        """Find node names from persistence directory."""
        adapters_dir = os.path.join(self.database_dir_, "adapters")
        return (
            list(os.listdir(adapters_dir))
            if os.path.exists(adapters_dir)
            else []
        )

    def find_json_importable_node_and_group_names(self):
        """Find node names and group names from persistence directory."""
        res = []
        adapters_dir = os.path.join(self.database_dir_, "adapters")
        if not os.path.exists(adapters_dir):
            return res

        for adapter_name in os.listdir(adapters_dir):
            groups_dir = os.path.join(adapters_dir, adapter_name, "groups")
            if not os.path.exists(groups_dir):
                continue
            for group_name in os.listdir(groups_dir):
                res.append((adapter_name, group_name))

        return res


if __name__ == "__main__":
    args = parse_args()

    VERBOSE = args.verbose

    migrations = Migration(schema_dir=args.schema_dir, database_dir=args.data_dir)

    if args.subcommand == "import":
        migrations.import_json()
    elif args.subcommand == "show":
        migrations.print_head()
    elif args.subcommand == "up":
        migrations.apply(args.import_once)