Schema migrations

class SchemaMigrations:
    """Manage the schema migrations in the database.
    It provides methods to create the schema_migrations table, check its existence,
    set the baseline version, and retrieve migration details.

    Version Added:
        1.0.0
    """

    def __init__(self, config: PumConfig) -> None:
        """Initialize the SchemaMigrations class with a database connection and configuration.

        Args:
            config (PumConfig): An instance of the PumConfig class containing configuration settings for the PUM system.

        """
        self.config = config
        self.migration_table_identifier = psycopg.sql.SQL(".").join(
            [
                psycopg.sql.Identifier(self.config.config.pum.migration_table_schema),
                psycopg.sql.Identifier(MIGRATION_TABLE_NAME),
            ]
        )
        self.migration_table_identifier_str = (
            f"{self.config.config.pum.migration_table_schema}.{MIGRATION_TABLE_NAME}"
        )

    def exists(self, connection: psycopg.Connection) -> bool:
        """Check if the schema_migrations information table exists.

        Args:
            connection: The database connection to check for the existence of the table.

        Returns:
            bool: True if the table exists, False otherwise.

        """
        query = psycopg.sql.SQL(
            """
        SELECT EXISTS (
            SELECT 1
            FROM information_schema.tables
            WHERE table_name = 'pum_migrations' AND table_schema = {schema}
        );
        """
        )

        parameters = {
            "schema": psycopg.sql.Literal(self.config.config.pum.migration_table_schema),
        }

        with connection.transaction():
            cursor = SqlContent(query).execute(connection, parameters=parameters)
            result = cursor._pum_results[0] if cursor._pum_results else None
            return result[0] if result else False

    @staticmethod
    def schemas_with_migrations(connection: psycopg.Connection) -> list[str]:
        """Check if the schema_migrations information table exists in any schemas.

        Args:
            connection: The database connection to check for the existence of the table.

        Returns:
            List[str]: List of all schemas where the table exists.

        """
        query = psycopg.sql.SQL(
            """
            SELECT table_schema
            FROM information_schema.tables
            WHERE table_name = 'pum_migrations'
        """
        )
        with connection.transaction():
            cursor = SqlContent(query).execute(connection, parameters={})
            return [row[0] for row in (cursor._pum_results or [])]

    @staticmethod
    def schemas_with_migration_details(
        connection: psycopg.Connection,
    ) -> list[dict]:
        """Return detailed migration info for each schema that has a pum_migrations table.

        For each schema, returns the module name, current version, first installation
        date and latest upgrade date (if different from the installation date).

        Args:
            connection: The database connection.

        Returns:
            list[dict]: A list of dicts with keys: schema, module, version,
                        installed_date, upgrade_date (None if never upgraded).

        Version Added:
            1.4.0
        """
        with connection.transaction():
            schemas = SchemaMigrations.schemas_with_migrations(connection)
            details = []
            for schema in schemas:
                table_id = psycopg.sql.SQL(".").join(
                    [psycopg.sql.Identifier(schema), psycopg.sql.Identifier(MIGRATION_TABLE_NAME)]
                )
                query = psycopg.sql.SQL(
                    """
                    SELECT
                        module,
                        version,
                        installed_date,
                        CASE WHEN upgrade_date > installed_date THEN upgrade_date END AS upgrade_date,
                        beta_testing
                    FROM (
                        SELECT
                            module,
                            (SELECT version FROM {table} ORDER BY version DESC, date_installed DESC LIMIT 1) AS version,
                            MIN(date_installed) AS installed_date,
                            MAX(date_installed) AS upgrade_date,
                            (SELECT beta_testing FROM {table} ORDER BY version DESC, date_installed DESC LIMIT 1) AS beta_testing
                        FROM {table}
                        GROUP BY module
                    ) sub
                    """
                )
                try:
                    cursor = SqlContent(query).execute(connection, parameters={"table": table_id})
                    for row in cursor._pum_results or []:
                        details.append(
                            {
                                "schema": schema,
                                "module": row[0],
                                "version": row[1],
                                "installed_date": row[2],
                                "upgrade_date": row[3],
                                "beta_testing": row[4],
                            }
                        )
                except Exception:
                    logger.warning(f"Could not read migration details from schema {schema}")
                    details.append(
                        {
                            "schema": schema,
                            "module": None,
                            "version": None,
                            "installed_date": None,
                            "upgrade_date": None,
                            "beta_testing": None,
                        }
                    )
            return details

    def exists_in_other_schemas(self, connection: psycopg.Connection) -> list[str]:
        """Check if the schema_migrations information table exists in other schemas.

        Args:
            connection: The database connection to check for the existence of the table.

        Returns:
            List[str]: List of schemas where the table exists.

        """
        with connection.transaction():
            all_schemas = self.schemas_with_migrations(connection)
            return [
                schema
                for schema in all_schemas
                if schema != self.config.config.pum.migration_table_schema
            ]

    def create(
        self,
        connection: psycopg.Connection,
        *,
        allow_multiple_modules: bool = False,
        commit: bool = False,
    ) -> None:
        """Create the schema_migrations information table
        Args:
            connection: The database connection to create the table.
            commit: If true, the transaction is committed. The default is false.
            allow_multiple_modules: If true, allows multiple PUM modules (with separate
                pum_migrations tables in different schemas) in the same database. Default is false.
        """
        if self.exists(connection):
            logger.debug(
                f"{self.config.config.pum.migration_table_schema}.pum_migrations table already exists."
            )
            return

        if not allow_multiple_modules:
            schemas = self.exists_in_other_schemas(connection)
            if len(schemas) > 1:
                raise PumSchemaMigrationError(
                    f"Another {MIGRATION_TABLE_NAME} table exists in other schemas {', '.join(schemas)}, indicating other PUM modules in this database. "
                    "Use allow_multiple_modules=True to allow multiple PUM modules per database."
                )
            elif len(schemas) == 1:
                raise PumSchemaMigrationError(
                    f"Another {MIGRATION_TABLE_NAME} table exists in schema {schemas[0]}, indicating another PUM module in this database. "
                    "Use allow_multiple_modules=True to allow multiple PUM modules per database."
                )
        # Create the schema if it doesn't exist
        parameters = {
            "version": psycopg.sql.Literal(MIGRATION_TABLE_VERSION),
            "schema": psycopg.sql.Identifier(self.config.config.pum.migration_table_schema),
            "table": self.migration_table_identifier,
        }

        create_schema_query = None
        if self.config.config.pum.migration_table_schema != "public":
            create_schema_query = psycopg.sql.SQL("CREATE SCHEMA IF NOT EXISTS {schema};")

        create_table_query = psycopg.sql.SQL(
            """CREATE TABLE IF NOT EXISTS {table}
            (
            id uuid DEFAULT gen_random_uuid() PRIMARY KEY,
            date_installed timestamp without time zone NOT NULL DEFAULT now(),
            module character varying(50) NOT NULL,
            version character varying(50) NOT NULL,
            beta_testing boolean NOT NULL DEFAULT false,
            changelog_files text[],
            parameters jsonb,
            migration_table_version integer NOT NULL DEFAULT {version}
            );
        """
        )

        comment_query = psycopg.sql.SQL("COMMENT ON TABLE {table} IS 'migration_table_version: 2';")

        if create_schema_query:
            SqlContent(create_schema_query).execute(connection, parameters=parameters)
        SqlContent(create_table_query).execute(connection, parameters=parameters)
        SqlContent(comment_query).execute(connection, parameters=parameters)

        logger.info(f"Created migration table: {self.migration_table_identifier_str}")

        if commit:
            connection.commit()

    def migration_table_version(self, connection: psycopg.Connection) -> int:
        """Return the migration table version.

        Args:
            connection: The database connection to get the migration table version.

        Returns:
            int | None: The migration table version, or None if the table does not exist.

        """
        query = psycopg.sql.SQL(
            """
            SELECT migration_table_version
            FROM {table}
            ORDER BY migration_table_version DESC
            LIMIT 1;
        """
        )

        parameters = {
            "table": self.migration_table_identifier,
        }

        with connection.transaction():
            cursor = SqlContent(query).execute(connection, parameters=parameters)
            row = cursor._pum_results[0] if cursor._pum_results else None
            if row is None:
                raise PumSchemaMigrationError(
                    f"Migration table {self.migration_table_identifier_str} does not exist."
                )
            return row[0]

    def update_migration_table_schema(self, connection: psycopg.Connection) -> None:
        """Update the migration table schema to the latest version.

        Args:
            connection: The database connection to update the table.

        """
        table_version = self.migration_table_version(connection)
        logger.info(
            f"Updating migration table {self.migration_table_identifier_str} from version {table_version} to {MIGRATION_TABLE_VERSION}."
        )
        if table_version == 1:
            alter_query = psycopg.sql.SQL("""
                                          ALTER TABLE {table} ALTER COLUMN migration_table_version ALTER TYPE integer SET DEFAULT {version} USING 1;
                                          ALTER TABLE {table} ALTER COLUMN module SET NOT NULL USING 'tww';
                                          """)
            parameters = {
                "table": self.migration_table_identifier,
                "version": psycopg.sql.Literal(MIGRATION_TABLE_VERSION),
            }
            SqlContent(alter_query).execute(connection, parameters=parameters)

    def set_baseline(
        self,
        connection: psycopg.Connection,
        version: packaging.version.Version | str,
        changelog_files: list[str] | None = None,
        parameters: dict | None = None,
        *,
        beta_testing: bool = False,
        commit: bool = False,
    ) -> None:
        """Set the baseline into the migration table.

        Args:
            connection: The database connection to set the baseline version.
            version: The version of the current database to set in the information.
            changelog_files: The list of changelog files that were applied.
            parameters: The parameters used in the migration.
            beta_testing: If true, the baseline is set to beta testing mode. The default is false.
            commit: If true, the transaction is committed. The default is False.

        """
        version_str = version
        version_packaging = version
        if isinstance(version_str, packaging.version.Version):
            version_str = str(version_str)
        if isinstance(version_packaging, str):
            version_packaging = packaging.version.parse(version_packaging)
        pattern = re.compile(r"^\d+\.\d+(\.\d+)?$")
        if not re.match(pattern, version_str):
            raise ValueError(f"Wrong version format: {version}. Must be x.y or x.y.z")

        try:
            current = self.baseline(connection=connection)
        except PumSchemaMigrationNoBaselineError:
            current = None
        if current:
            self.update_migration_table_schema(connection)
        if current and current >= version_packaging:
            raise PumSchemaMigrationError(
                f"Cannot set baseline {version_str} as it is already set at {current}."
            )

        code = psycopg.sql.SQL("""
INSERT INTO {table} (
    module,
    version,
    beta_testing,
    migration_table_version,
    changelog_files,
    parameters
) VALUES (
    {module},
    {version},
    {beta_testing},
    {migration_table_version},
    {changelog_files},
    {parameters}
);""")

        query_parameters = {
            "table": self.migration_table_identifier,
            "module": psycopg.sql.Literal(self.config.config.pum.module),
            "version": psycopg.sql.Literal(version_str),
            "beta_testing": psycopg.sql.Literal(beta_testing),
            "migration_table_version": psycopg.sql.Literal(MIGRATION_TABLE_VERSION),
            "changelog_files": psycopg.sql.Literal(changelog_files or []),
            "parameters": psycopg.sql.Literal(json.dumps(parameters or {})),
        }

        logger.info(
            f"Setting baseline version {version} in {self.migration_table_identifier_str} table"
        )
        SqlContent(code).execute(connection, parameters=query_parameters, commit=commit)

    def has_baseline(self, connection: psycopg.Connection) -> bool:
        """Check if the migration table has a baseline version.

        Args:
            connection: The database connection to check for the baseline version.
        Returns:
            bool: True if the baseline version exists, False otherwise.
        """
        with connection.transaction():
            try:
                self.baseline(connection=connection)
                return True
            except PumSchemaMigrationError:
                return False

    def baseline(self, connection: psycopg.Connection) -> packaging.version.Version:
        """Return the baseline version from the migration table.

        Args:
            connection: psycopg.Connection
                The database connection to get the baseline version.

        Returns:
            packaging.version.Version | None: The baseline version.

        Raises:
            PumSchemaMigrationError: If the migration table does not exist or if no baseline version is found
            PumSchemaMigrationNoBaselineError: If the migration table does not exist
        """

        with connection.transaction():
            if not self.exists(connection=connection):
                raise PumSchemaMigrationError(
                    f"{self.migration_table_identifier_str} table does not exist."
                )

            query = psycopg.sql.SQL(
                """
                SELECT version
                FROM {table}
                WHERE id = (
                    SELECT id
                    FROM {table}
                    ORDER BY version DESC, date_installed DESC
                    LIMIT 1
                )
            """
            )

            parameters = {
                "table": self.migration_table_identifier,
            }

            cursor = SqlContent(query).execute(connection, parameters=parameters)
            row = cursor._pum_results[0] if cursor._pum_results else None
            if row is None:
                raise PumSchemaMigrationNoBaselineError(
                    f"Baseline version not found in the {self.migration_table_identifier_str} table."
                )
            return packaging.version.parse(row[0])

    def migration_summary(self, connection: psycopg.Connection) -> dict:
        """Return aggregated migration summary for the configured schema.

        Returns the module name, current version, first installation date,
        latest upgrade date (if different from the installation date),
        beta testing flag, and parameters.

        Args:
            connection: The database connection.

        Returns:
            dict: A dict with keys: schema, module, version, installed_date,
                  upgrade_date (None if never upgraded), beta_testing, parameters.

        Raises:
            PumSchemaMigrationError: If the migration table does not exist or has no data.
        """
        query = psycopg.sql.SQL(
            """
            SELECT
                module,
                version,
                installed_date,
                CASE WHEN upgrade_date > installed_date THEN upgrade_date END AS upgrade_date,
                beta_testing,
                parameters
            FROM (
                SELECT
                    module,
                    (SELECT version FROM {table} ORDER BY version DESC, date_installed DESC LIMIT 1) AS version,
                    MIN(date_installed) AS installed_date,
                    MAX(date_installed) AS upgrade_date,
                    (SELECT beta_testing FROM {table} ORDER BY version DESC, date_installed DESC LIMIT 1) AS beta_testing,
                    (SELECT parameters FROM {table} ORDER BY version DESC, date_installed DESC LIMIT 1) AS parameters
                FROM {table}
                GROUP BY module
            ) sub
            """
        )
        with connection.transaction():
            cursor = SqlContent(query).execute(
                connection, parameters={"table": self.migration_table_identifier}
            )
            row = cursor._pum_results[0] if cursor._pum_results else None
            if row is None:
                raise PumSchemaMigrationError(
                    f"No migration data found in {self.migration_table_identifier_str}."
                )
            return {
                "schema": self.config.config.pum.migration_table_schema,
                "module": row[0],
                "version": row[1],
                "installed_date": row[2],
                "upgrade_date": row[3],
                "beta_testing": row[4],
                "parameters": row[5],
            }

    def migration_details(self, connection: psycopg.Connection, version: str | None = None) -> dict:
        """Return the migration details from the migration table.

        Args:
            connection:
                The database connection to get the migration details.
            version:
                The version of the migration to get details for.
                If None, last migration is returned.

        Returns:
            dict: The migration details.

        Raises:
            PumSchemaMigrationError: If the migration table does not exist or if no migration details are found.
        """
        query = None
        if version is None:
            query = psycopg.sql.SQL(
                """
                SELECT *
                FROM {table}
                WHERE id = (
                        SELECT id
                        FROM {table}
                        ORDER BY version DESC, date_installed DESC
                        LIMIT 1
                    )
                ORDER BY date_installed DESC
            """
            )

            parameters = {
                "table": self.migration_table_identifier,
            }
        else:
            query = psycopg.sql.SQL(
                """
                SELECT *
                FROM {table}
                WHERE version = {version}
            """
            )

            parameters = {
                "table": self.migration_table_identifier,
                "version": psycopg.sql.Literal(version),
            }

        with connection.transaction():
            cursor = SqlContent(query).execute(connection, parameters=parameters)
            row = cursor._pum_results[0] if cursor._pum_results else None
            if row is None:
                raise PumSchemaMigrationError(
                    f"Migration details not found for version {version} in the {self.migration_table_identifier_str} table."
                )
            return dict(zip([desc[0] for desc in cursor._pum_description], row, strict=False))

    def compare(self, connection: psycopg.Connection) -> int:
        """Compare the migrations details in the database to the changelogs in the source.

        Args:
            connection: The database connection to get the baseline version.
        Returns:
            int: -1 if database is behind, 0 if up to date.

        Raises:
            PumSchemaMigrationError: If there is a mismatch between the database and the source.
        """

        with connection.transaction():
            current_version = self.baseline(connection=connection)
            migration_details = self.migration_details(connection=connection)
            changelogs = [str(changelog.version) for changelog in self.config.changelogs()]

            # Check if the current migration version is in the changelogs
            if migration_details["version"] not in changelogs:
                raise PumSchemaMigrationError(
                    f"Changelog for version {migration_details['version']} not found in the source."
                )

            # Check if there are newer changelogs than current version
            for changelog_version in changelogs:
                if packaging.version.parse(changelog_version) > current_version:
                    return -1  # database is behind

            return 0  # database is up to date