diff --git a/alembic.ini b/tdp/alembic.ini similarity index 100% rename from alembic.ini rename to tdp/alembic.ini diff --git a/alembic_migration/README.md b/tdp/alembic_migration/README.md similarity index 100% rename from alembic_migration/README.md rename to tdp/alembic_migration/README.md diff --git a/alembic_migration/mysql/env.py b/tdp/alembic_migration/mysql/env.py similarity index 100% rename from alembic_migration/mysql/env.py rename to tdp/alembic_migration/mysql/env.py diff --git a/alembic_migration/mysql/script.py.mako b/tdp/alembic_migration/mysql/script.py.mako similarity index 100% rename from alembic_migration/mysql/script.py.mako rename to tdp/alembic_migration/mysql/script.py.mako diff --git a/alembic_migration/mysql/versions/.keep b/tdp/alembic_migration/mysql/versions/.keep similarity index 100% rename from alembic_migration/mysql/versions/.keep rename to tdp/alembic_migration/mysql/versions/.keep diff --git a/tdp/alembic_migration/mysql/versions/tdp_lib_1.1_initial_table_creation.py b/tdp/alembic_migration/mysql/versions/tdp_lib_1.1_initial_table_creation.py new file mode 100644 index 00000000..9881387b --- /dev/null +++ b/tdp/alembic_migration/mysql/versions/tdp_lib_1.1_initial_table_creation.py @@ -0,0 +1,108 @@ +# Copyright 2022 TOSIT.IO +# SPDX-License-Identifier: Apache-2.0 + +"""Initial table creation + +Revision ID: tdp_lib_1.1 +Revises: +Create Date: 2024-06-11 09:59:57.688453 + +""" +from typing import Sequence, Union + +import sqlalchemy as sa +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "tdp_lib_1.1" +down_revision: Union[str, None] = None +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.create_table( + "deployment", + sa.Column("id", sa.Integer(), nullable=False), + sa.Column("options", sa.JSON(none_as_null=True), nullable=True), + sa.Column("start_time", sa.DateTime(), nullable=True), + sa.Column("end_time", sa.DateTime(), nullable=True), + sa.Column( + "state", + sa.Enum( + "PLANNED", "RUNNING", "SUCCESS", "FAILURE", name="deploymentstateenum" + ), + nullable=True, + ), + sa.Column( + "deployment_type", + sa.Enum( + "DAG", + "OPERATIONS", + "RESUME", + "RECONFIGURE", + "CUSTOM", + name="deploymenttypeenum", + ), + nullable=True, + ), + sa.PrimaryKeyConstraint("id"), + ) + op.create_table( + "operation", + sa.Column("deployment_id", sa.Integer(), nullable=False), + sa.Column("operation_order", sa.Integer(), nullable=False), + sa.Column("operation", sa.String(length=72), nullable=False), + sa.Column("host", sa.String(length=255), nullable=True), + sa.Column("extra_vars", sa.JSON(none_as_null=True), nullable=True), + sa.Column("start_time", sa.DateTime(), nullable=True), + sa.Column("end_time", sa.DateTime(), nullable=True), + sa.Column( + "state", + sa.Enum( + "PLANNED", + "RUNNING", + "PENDING", + "SUCCESS", + "FAILURE", + "HELD", + name="operationstateenum", + ), + nullable=False, + ), + sa.Column("logs", sa.LargeBinary(length=10000000), nullable=True), + sa.ForeignKeyConstraint( + ["deployment_id"], + ["deployment.id"], + ), + sa.PrimaryKeyConstraint("deployment_id", "operation_order"), + ) + op.create_table( + "sch_status_log", + sa.Column("id", sa.Integer(), nullable=False), + sa.Column("event_time", sa.DateTime(), nullable=False), + sa.Column("service", sa.String(length=20), nullable=False), + sa.Column("component", sa.String(length=30), nullable=True), + sa.Column("host", sa.String(length=255), nullable=True), + sa.Column("running_version", sa.String(length=40), nullable=True), + sa.Column("configured_version", sa.String(length=40), nullable=True), + sa.Column("to_config", sa.Boolean(), nullable=True), + sa.Column("to_restart", sa.Boolean(), nullable=True), + sa.Column("is_active", sa.Boolean(), nullable=True), + sa.Column( + "source", + sa.Enum( + "DEPLOYMENT", "FORCED", "STALE", "MANUAL", name="schstatuslogsourceenum" + ), + nullable=False, + ), + sa.Column("deployment_id", sa.Integer(), nullable=True), + sa.Column("message", sa.String(length=512), nullable=True), + sa.ForeignKeyConstraint( + ["deployment_id"], + ["deployment.id"], + ), + sa.PrimaryKeyConstraint("id"), + ) + # ### end Alembic commands ### diff --git a/alembic_migration/postgresql/env.py b/tdp/alembic_migration/postgresql/env.py similarity index 100% rename from alembic_migration/postgresql/env.py rename to tdp/alembic_migration/postgresql/env.py diff --git a/alembic_migration/postgresql/script.py.mako b/tdp/alembic_migration/postgresql/script.py.mako similarity index 100% rename from alembic_migration/postgresql/script.py.mako rename to tdp/alembic_migration/postgresql/script.py.mako diff --git a/alembic_migration/postgresql/versions/.keep b/tdp/alembic_migration/postgresql/versions/.keep similarity index 100% rename from alembic_migration/postgresql/versions/.keep rename to tdp/alembic_migration/postgresql/versions/.keep diff --git a/tdp/alembic_migration/postgresql/versions/tdp_lib_1.1_initial_table_creation.py b/tdp/alembic_migration/postgresql/versions/tdp_lib_1.1_initial_table_creation.py new file mode 100644 index 00000000..d4423d4f --- /dev/null +++ b/tdp/alembic_migration/postgresql/versions/tdp_lib_1.1_initial_table_creation.py @@ -0,0 +1,144 @@ +# Copyright 2022 TOSIT.IO +# SPDX-License-Identifier: Apache-2.0 + +"""Initial table creation + +Revision ID: tdp_lib_1.1 +Revises: +Create Date: 2024-06-11 10:00:08.925442 + +""" +from typing import Sequence, Union + +import sqlalchemy as sa +from alembic import op +from sqlalchemy.dialects import postgresql + +# revision identifiers, used by Alembic. +revision: str = "tdp_lib_1.1" +down_revision: Union[str, None] = None +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + sa.Enum( + "DEPLOYMENT", "FORCED", "STALE", "MANUAL", name="schstatuslogsourceenum" + ).create(op.get_bind()) + sa.Enum( + "DAG", + "OPERATIONS", + "RESUME", + "RECONFIGURE", + "CUSTOM", + name="deploymenttypeenum", + ).create(op.get_bind()) + sa.Enum( + "PLANNED", "RUNNING", "SUCCESS", "FAILURE", name="deploymentstateenum" + ).create(op.get_bind()) + sa.Enum( + "PLANNED", + "RUNNING", + "PENDING", + "SUCCESS", + "FAILURE", + "HELD", + name="operationstateenum", + ).create(op.get_bind()) + op.create_table( + "deployment", + sa.Column("id", sa.Integer(), nullable=False), + sa.Column("options", sa.JSON(none_as_null=True), nullable=True), + sa.Column("start_time", sa.DateTime(), nullable=True), + sa.Column("end_time", sa.DateTime(), nullable=True), + sa.Column( + "state", + postgresql.ENUM( + "PLANNED", + "RUNNING", + "SUCCESS", + "FAILURE", + name="deploymentstateenum", + create_type=False, + ), + nullable=True, + ), + sa.Column( + "deployment_type", + postgresql.ENUM( + "DAG", + "OPERATIONS", + "RESUME", + "RECONFIGURE", + "CUSTOM", + name="deploymenttypeenum", + create_type=False, + ), + nullable=True, + ), + sa.PrimaryKeyConstraint("id"), + ) + op.create_table( + "operation", + sa.Column("deployment_id", sa.Integer(), nullable=False), + sa.Column("operation_order", sa.Integer(), nullable=False), + sa.Column("operation", sa.String(length=72), nullable=False), + sa.Column("host", sa.String(length=255), nullable=True), + sa.Column("extra_vars", sa.JSON(none_as_null=True), nullable=True), + sa.Column("start_time", sa.DateTime(), nullable=True), + sa.Column("end_time", sa.DateTime(), nullable=True), + sa.Column( + "state", + postgresql.ENUM( + "PLANNED", + "RUNNING", + "PENDING", + "SUCCESS", + "FAILURE", + "HELD", + name="operationstateenum", + create_type=False, + ), + nullable=False, + ), + sa.Column("logs", sa.LargeBinary(length=10000000), nullable=True), + sa.ForeignKeyConstraint( + ["deployment_id"], + ["deployment.id"], + ), + sa.PrimaryKeyConstraint("deployment_id", "operation_order"), + ) + op.create_table( + "sch_status_log", + sa.Column("id", sa.Integer(), nullable=False), + sa.Column("event_time", sa.DateTime(), nullable=False), + sa.Column("service", sa.String(length=20), nullable=False), + sa.Column("component", sa.String(length=30), nullable=True), + sa.Column("host", sa.String(length=255), nullable=True), + sa.Column("running_version", sa.String(length=40), nullable=True), + sa.Column("configured_version", sa.String(length=40), nullable=True), + sa.Column("to_config", sa.Boolean(), nullable=True), + sa.Column("to_restart", sa.Boolean(), nullable=True), + sa.Column("is_active", sa.Boolean(), nullable=True), + sa.Column( + "source", + postgresql.ENUM( + "DEPLOYMENT", + "FORCED", + "STALE", + "MANUAL", + name="schstatuslogsourceenum", + create_type=False, + ), + nullable=False, + ), + sa.Column("deployment_id", sa.Integer(), nullable=True), + sa.Column("message", sa.String(length=512), nullable=True), + sa.ForeignKeyConstraint( + ["deployment_id"], + ["deployment.id"], + ), + sa.PrimaryKeyConstraint("id"), + ) + # ### end Alembic commands ### diff --git a/alembic_migration/sqlite/env.py b/tdp/alembic_migration/sqlite/env.py similarity index 100% rename from alembic_migration/sqlite/env.py rename to tdp/alembic_migration/sqlite/env.py diff --git a/alembic_migration/sqlite/script.py.mako b/tdp/alembic_migration/sqlite/script.py.mako similarity index 100% rename from alembic_migration/sqlite/script.py.mako rename to tdp/alembic_migration/sqlite/script.py.mako diff --git a/alembic_migration/sqlite/versions/.keep b/tdp/alembic_migration/sqlite/versions/.keep similarity index 100% rename from alembic_migration/sqlite/versions/.keep rename to tdp/alembic_migration/sqlite/versions/.keep diff --git a/tdp/alembic_migration/sqlite/versions/tdp_lib_1.1_initial_table_creation.py b/tdp/alembic_migration/sqlite/versions/tdp_lib_1.1_initial_table_creation.py new file mode 100644 index 00000000..c53564a3 --- /dev/null +++ b/tdp/alembic_migration/sqlite/versions/tdp_lib_1.1_initial_table_creation.py @@ -0,0 +1,108 @@ +# Copyright 2022 TOSIT.IO +# SPDX-License-Identifier: Apache-2.0 + +"""Initial table creation + +Revision ID: tdp_lib_1.1 +Revises: +Create Date: 2024-06-11 09:59:29.908273 + +""" +from typing import Sequence, Union + +import sqlalchemy as sa +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "tdp_lib_1.1" +down_revision: Union[str, None] = None +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.create_table( + "deployment", + sa.Column("id", sa.Integer(), nullable=False), + sa.Column("options", sa.JSON(none_as_null=True), nullable=True), + sa.Column("start_time", sa.DateTime(), nullable=True), + sa.Column("end_time", sa.DateTime(), nullable=True), + sa.Column( + "state", + sa.Enum( + "PLANNED", "RUNNING", "SUCCESS", "FAILURE", name="deploymentstateenum" + ), + nullable=True, + ), + sa.Column( + "deployment_type", + sa.Enum( + "DAG", + "OPERATIONS", + "RESUME", + "RECONFIGURE", + "CUSTOM", + name="deploymenttypeenum", + ), + nullable=True, + ), + sa.PrimaryKeyConstraint("id"), + ) + op.create_table( + "operation", + sa.Column("deployment_id", sa.Integer(), nullable=False), + sa.Column("operation_order", sa.Integer(), nullable=False), + sa.Column("operation", sa.String(length=72), nullable=False), + sa.Column("host", sa.String(length=255), nullable=True), + sa.Column("extra_vars", sa.JSON(none_as_null=True), nullable=True), + sa.Column("start_time", sa.DateTime(), nullable=True), + sa.Column("end_time", sa.DateTime(), nullable=True), + sa.Column( + "state", + sa.Enum( + "PLANNED", + "RUNNING", + "PENDING", + "SUCCESS", + "FAILURE", + "HELD", + name="operationstateenum", + ), + nullable=False, + ), + sa.Column("logs", sa.LargeBinary(length=10000000), nullable=True), + sa.ForeignKeyConstraint( + ["deployment_id"], + ["deployment.id"], + ), + sa.PrimaryKeyConstraint("deployment_id", "operation_order"), + ) + op.create_table( + "sch_status_log", + sa.Column("id", sa.Integer(), nullable=False), + sa.Column("event_time", sa.DateTime(), nullable=False), + sa.Column("service", sa.String(length=20), nullable=False), + sa.Column("component", sa.String(length=30), nullable=True), + sa.Column("host", sa.String(length=255), nullable=True), + sa.Column("running_version", sa.String(length=40), nullable=True), + sa.Column("configured_version", sa.String(length=40), nullable=True), + sa.Column("to_config", sa.Boolean(), nullable=True), + sa.Column("to_restart", sa.Boolean(), nullable=True), + sa.Column("is_active", sa.Boolean(), nullable=True), + sa.Column( + "source", + sa.Enum( + "DEPLOYMENT", "FORCED", "STALE", "MANUAL", name="schstatuslogsourceenum" + ), + nullable=False, + ), + sa.Column("deployment_id", sa.Integer(), nullable=True), + sa.Column("message", sa.String(length=512), nullable=True), + sa.ForeignKeyConstraint( + ["deployment_id"], + ["deployment.id"], + ), + sa.PrimaryKeyConstraint("id"), + ) + # ### end Alembic commands ### diff --git a/tdp/cli/commands/init.py b/tdp/cli/commands/init.py index dc52c46f..32e1d176 100644 --- a/tdp/cli/commands/init.py +++ b/tdp/cli/commands/init.py @@ -3,11 +3,14 @@ from __future__ import annotations +import logging from pathlib import Path from typing import TYPE_CHECKING import click -from sqlalchemy import Engine +from alembic.command import stamp, upgrade +from alembic.config import Config +from sqlalchemy import MetaData, Table, create_engine from tdp.cli.params import ( collections_option, @@ -17,10 +20,13 @@ ) from tdp.core.models import init_database from tdp.core.variables import ClusterVariables +from tdp.find_tdp_lib_root_folder import find_tdp_lib_root_folder if TYPE_CHECKING: from tdp.core.collections import Collections +logging.getLogger("alembic").setLevel(logging.ERROR) + @click.command() @click.option( @@ -32,18 +38,45 @@ help="Path to TDP variables overrides. Can be used multiple times. Last one takes precedence.", ) @collections_option -@database_dsn_option +@database_dsn_option(create_engine=False) @validate_option @vars_option(exists=False) def init( overrides: tuple[Path], collections: Collections, - db_engine: Engine, + database_dsn: str, validate: bool, vars: Path, ): """Initialize the database and the TDP variables.""" - init_database(db_engine) + engine = create_engine(database_dsn) + with engine.connect() as connection: + tdp_lib_folder_path = find_tdp_lib_root_folder() + if tdp_lib_folder_path and not (tdp_lib_folder_path / "alembic.ini").exists(): + raise click.ClickException("alembic.ini file could not be found.") + + alembic_config = Config( + file_=f"{tdp_lib_folder_path}/alembic.ini", + ini_section=engine.dialect.name, + ) + alembic_config.set_main_option("sqlalchemy.url", database_dsn) + + try: + alembic_version = Table("alembic_version", MetaData(), autoload_with=engine) + db_revision_id_row = connection.execute(alembic_version.select()).fetchone() + db_revision_id = db_revision_id_row[0] if db_revision_id_row else None + + except: + db_revision_id = None + + # Upgrade database if new revision in migration folder, else create all tables or fail if revisions don't coincide. + if db_revision_id: + upgrade(config=alembic_config, revision="head") + click.echo(f"Upgrade tables to revision ID: {db_revision_id}") + else: + init_database(engine) + stamp(config=alembic_config, revision="head") + # Create the cluster variables. cluster_variables = ClusterVariables.initialize_cluster_variables( collections, vars, overrides, validate=validate ) diff --git a/tdp/cli/params/database_dsn_option.py b/tdp/cli/params/database_dsn_option.py index 29385d8e..c3b2d9be 100644 --- a/tdp/cli/params/database_dsn_option.py +++ b/tdp/cli/params/database_dsn_option.py @@ -33,7 +33,7 @@ def database_dsn_option( def decorator(fn: FC) -> FC: return click.option( - "db_engine", + "db_engine" if create_engine else "database_dsn", "--database-dsn", envvar="TDP_DATABASE_DSN", required=True, diff --git a/tdp/find_tdp_lib_root_folder.py b/tdp/find_tdp_lib_root_folder.py new file mode 100644 index 00000000..6478523d --- /dev/null +++ b/tdp/find_tdp_lib_root_folder.py @@ -0,0 +1,27 @@ +# Copyright 2022 TOSIT.IO +# SPDX-License-Identifier: Apache-2.0 + +from pathlib import Path +from typing import Optional + + +def find_tdp_lib_root_folder(current_path: Optional[Path] = None): + # If current_path is not provided, start from the directory of this file + if current_path is None: + current_path = Path(__file__).parent + + while True: + # Check if 'pyproject.toml' exists in the current directory + if (current_path / "alembic.ini").exists(): + return current_path.resolve() + + parent_path = current_path.parent + # If the parent path is the same as the current path, we have reached the root directory + if parent_path == current_path: + return None + + current_path = parent_path + + +if __name__ == "__main__": + print(find_tdp_lib_root_folder()) diff --git a/scripts/__init__.py b/tdp/scripts/__init__.py similarity index 100% rename from scripts/__init__.py rename to tdp/scripts/__init__.py diff --git a/scripts/playbooks.py b/tdp/scripts/playbooks.py similarity index 100% rename from scripts/playbooks.py rename to tdp/scripts/playbooks.py diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 9f3cc909..1c369357 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -7,7 +7,7 @@ import pytest from click.testing import CliRunner -from sqlalchemy import create_engine +from sqlalchemy import MetaData, Table, create_engine from tdp.cli.commands.init import init from tdp.core.models import BaseModel @@ -37,6 +37,7 @@ def tdp_init( yield TDPInitArgs(collection_path, db_dsn, vars) engine = create_engine(db_dsn) BaseModel.metadata.drop_all(engine) + Table("alembic_version", MetaData(), autoload_with=engine).drop(engine) engine.dispose()