pum 1.2.2__py3-none-any.whl → 1.3.0__py3-none-any.whl

pum/role_manager.py

@@ -1,10 +1,14 @@
 import enum
-from typing import Optional
+from typing import Optional, TYPE_CHECKING
 import copy
 import psycopg
 import logging
 
 from .sql_content import SqlContent
+from .exceptions import PumException
+
+if TYPE_CHECKING:
+    from .feedback import Feedback
 
 logger = logging.getLogger(__name__)
 
@@ -40,12 +44,14 @@ class Permission:
         role: str,
         connection: psycopg.Connection,
         commit: bool = False,
+        feedback: Optional["Feedback"] = None,
     ) -> None:
         """Grant the permission to the specified role.
         Args:
             role: The name of the role to grant the permission to.
             connection: The database connection to execute the SQL statements.
             commit: Whether to commit the transaction. Defaults to False.
+            feedback: Optional feedback object for progress reporting.
         """
         if not isinstance(role, str):
             raise TypeError("Role must be a string.")
@@ -54,13 +60,37 @@ class Permission:
             raise ValueError("Schemas must be defined for the permission.")
 
         for schema in self.schemas:
-            logger.info(f"Granting {self.type.value} permission on schema {schema} to role {role}.")
+            if feedback and feedback.is_cancelled():
+                raise PumException("Permission grant cancelled by user")
+
+            # Detect if schema exists; if not, warn and continue
+            cursor = SqlContent("SELECT 1 FROM pg_namespace WHERE nspname = {schema}").execute(
+                connection=connection,
+                commit=False,
+                parameters={"schema": psycopg.sql.Literal(schema)},
+            )
+            if not cursor._pum_results or not cursor._pum_results[0]:
+                logger.warning(
+                    f"Schema {schema} does not exist; skipping grant of {self.type.value} "
+                    f"permission to role {role}."
+                )
+                continue
+
+            logger.debug(
+                f"Granting {self.type.value} permission on schema {schema} to role {role}."
+            )
             if self.type == PermissionType.READ:
                 SqlContent("""
                         GRANT USAGE ON SCHEMA {schema} TO {role};
-                        GRANT USAGE, SELECT ON ALL SEQUENCES IN SCHEMA {schema} TO {role};
                         GRANT SELECT, REFERENCES, TRIGGER ON ALL TABLES IN SCHEMA {schema} TO {role};
+                        GRANT USAGE, SELECT ON ALL SEQUENCES IN SCHEMA {schema} TO {role};
+                        GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA {schema} TO {role};
+                        GRANT EXECUTE ON ALL ROUTINES IN SCHEMA {schema} TO {role};
                         ALTER DEFAULT PRIVILEGES IN SCHEMA {schema} GRANT SELECT, REFERENCES, TRIGGER ON TABLES TO {role};
+                        ALTER DEFAULT PRIVILEGES IN SCHEMA {schema} GRANT SELECT ON SEQUENCES TO {role};
+                        ALTER DEFAULT PRIVILEGES IN SCHEMA {schema} GRANT EXECUTE ON FUNCTIONS TO {role};
+                        ALTER DEFAULT PRIVILEGES IN SCHEMA {schema} GRANT EXECUTE ON ROUTINES TO {role};
+                        ALTER DEFAULT PRIVILEGES IN SCHEMA {schema} GRANT USAGE ON TYPES TO {role};
                            """).execute(
                     connection=connection,
                     commit=False,
@@ -69,12 +99,20 @@ class Permission:
                         "role": psycopg.sql.Identifier(role),
                     },
                 )
+                # Grant permissions on existing types
+                self._grant_existing_types(connection, schema, role, "USAGE")
             elif self.type == PermissionType.WRITE:
                 SqlContent("""
                         GRANT ALL ON SCHEMA {schema} TO {role};
                         GRANT ALL ON ALL TABLES IN SCHEMA {schema} TO {role};
                         GRANT ALL ON ALL SEQUENCES IN SCHEMA {schema} TO {role};
+                        GRANT ALL ON ALL FUNCTIONS IN SCHEMA {schema} TO {role};
+                        GRANT ALL ON ALL ROUTINES IN SCHEMA {schema} TO {role};
                         ALTER DEFAULT PRIVILEGES IN SCHEMA {schema} GRANT ALL ON TABLES TO {role};
+                        ALTER DEFAULT PRIVILEGES IN SCHEMA {schema} GRANT ALL ON SEQUENCES TO {role};
+                        ALTER DEFAULT PRIVILEGES IN SCHEMA {schema} GRANT ALL ON FUNCTIONS TO {role};
+                        ALTER DEFAULT PRIVILEGES IN SCHEMA {schema} GRANT ALL ON ROUTINES TO {role};
+                        ALTER DEFAULT PRIVILEGES IN SCHEMA {schema} GRANT ALL ON TYPES TO {role};
                            """).execute(
                     connection=connection,
                     commit=False,
@@ -83,15 +121,62 @@ class Permission:
                         "role": psycopg.sql.Identifier(role),
                     },
                 )
+                # Grant permissions on existing types
+                self._grant_existing_types(connection, schema, role, "ALL")
             else:
                 raise ValueError(f"Unknown permission type: {self.type}")
 
         if commit:
+            if feedback:
+                feedback.lock_cancellation()
             connection.commit()
 
+    def _grant_existing_types(
+        self, connection: psycopg.Connection, schema: str, role: str, privilege: str
+    ) -> None:
+        """Grant permissions on all existing types in a schema.
+
+        Args:
+            connection: The database connection.
+            schema: The schema name.
+            role: The role name.
+            privilege: The privilege to grant (USAGE or ALL).
+
+        """
+        # Query for all types in the schema (excluding array types and internal types)
+        cursor = connection.cursor()
+        cursor.execute(
+            """
+            SELECT t.typname
+            FROM pg_type t
+            JOIN pg_namespace n ON t.typnamespace = n.oid
+            WHERE n.nspname = %s
+              AND t.typtype IN ('e', 'c', 'd', 'b', 'r')  -- enum, composite, domain, base, range
+              AND t.typname NOT LIKE '_%%'  -- exclude array types
+            """,
+            (schema,),
+        )
+        types = cursor.fetchall()
+
+        # Grant permissions on each type
+        for (type_name,) in types:
+            grant_sql = psycopg.sql.SQL(
+                "GRANT {privilege} ON TYPE {schema}.{type_name} TO {role}"
+            ).format(
+                privilege=psycopg.sql.SQL(privilege),
+                schema=psycopg.sql.Identifier(schema),
+                type_name=psycopg.sql.Identifier(type_name),
+                role=psycopg.sql.Identifier(role),
+            )
+            cursor.execute(grant_sql)
+
+    def __repr__(self) -> str:
+        """Return a string representation of the Permission object."""
+        return f"<Permission: {self.type.value} on {self.schemas}>"
+
 
 class Role:
-    """ "
+    """
     Represents a database role with associated permissions and optional inheritance.
     The Role class encapsulates the concept of a database role, including its name,
     permissions, optional inheritance from another role, and an optional description.
@@ -138,30 +223,36 @@ class Role:
         Returns:
             bool: True if the role exists, False otherwise.
         """
-        return (
-            SqlContent("SELECT 1 FROM pg_roles WHERE rolname = {name}")
-            .execute(
-                connection=connection,
-                commit=False,
-                parameters={"name": psycopg.sql.Literal(self.name)},
-            )
-            .fetchone()
-            is not None
+        cursor = SqlContent("SELECT 1 FROM pg_roles WHERE rolname = {name}").execute(
+            connection=connection,
+            commit=False,
+            parameters={"name": psycopg.sql.Literal(self.name)},
         )
+        return bool(cursor._pum_results)
 
     def create(
-        self, connection: psycopg.Connection, grant: bool = False, commit: bool = False
+        self,
+        connection: psycopg.Connection,
+        grant: bool = False,
+        commit: bool = False,
+        feedback: Optional["Feedback"] = None,
     ) -> None:
         """Create the role in the database.
         Args:
             connection: The database connection to execute the SQL statements.
             grant: Whether to grant permissions to the role. Defaults to False.
             commit: Whether to commit the transaction. Defaults to False.
+            feedback: Optional feedback object for progress reporting.
         """
+        if feedback and feedback.is_cancelled():
+            from .exceptions import PumException
+
+            raise PumException("Role creation cancelled by user")
+
         if self.exists(connection):
-            logger.info(f"Role {self.name} already exists, skipping creation.")
+            logger.debug(f"Role {self.name} already exists, skipping creation.")
         else:
-            logger.info(f"Creating role {self.name}.")
+            logger.debug(f"Creating role {self.name}.")
             SqlContent(
                 "CREATE ROLE {name} NOSUPERUSER INHERIT NOCREATEDB NOCREATEROLE NOREPLICATION"
             ).execute(
@@ -189,9 +280,13 @@ class Role:
                 )
         if grant:
             for permission in self.permissions():
-                permission.grant(role=self.name, connection=connection, commit=commit)
+                permission.grant(
+                    role=self.name, connection=connection, commit=False, feedback=feedback
+                )
 
         if commit:
+            if feedback:
+                feedback.lock_cancellation()
             connection.commit()
 
 
@@ -232,28 +327,61 @@ class RoleManager:
                 )
 
     def create_roles(
-        self, connection: psycopg.Connection, grant: bool = False, commit: bool = False
+        self,
+        connection: psycopg.Connection,
+        grant: bool = False,
+        commit: bool = False,
+        feedback: Optional["Feedback"] = None,
     ) -> None:
         """Create roles in the database.
         Args:
             connection: The database connection to execute the SQL statements.
             grant: Whether to grant permissions to the roles. Defaults to False.
             commit: Whether to commit the transaction. Defaults to False.
+            feedback: Optional feedback object for progress reporting.
         """
-        for role in self.roles.values():
-            role.create(connection=connection, commit=False, grant=grant)
+        roles_list = list(self.roles.values())
+        for role in roles_list:
+            if feedback and feedback.is_cancelled():
+                from .exceptions import PumException
+
+                raise PumException("Role creation cancelled by user")
+            if feedback:
+                feedback.increment_step()
+                feedback.report_progress(f"Creating role: {role.name}")
+            role.create(connection=connection, commit=False, grant=grant, feedback=feedback)
         if commit:
+            if feedback:
+                feedback.lock_cancellation()
             connection.commit()
 
-    def grant_permissions(self, connection: psycopg.Connection, commit: bool = False) -> None:
+    def grant_permissions(
+        self,
+        connection: psycopg.Connection,
+        commit: bool = False,
+        feedback: Optional["Feedback"] = None,
+    ) -> None:
         """Grant permissions to the roles in the database.
         Args:
             connection: The database connection to execute the SQL statements.
             commit: Whether to commit the transaction. Defaults to False.
+            feedback: Optional feedback object for progress reporting.
         """
-        for role in self.roles.values():
+        roles_list = list(self.roles.values())
+        for role in roles_list:
+            if feedback and feedback.is_cancelled():
+                from .exceptions import PumException
+
+                raise PumException("Permission grant cancelled by user")
+            if feedback:
+                feedback.increment_step()
+                feedback.report_progress(f"Granting permissions to role: {role.name}")
             for permission in role.permissions():
-                permission.grant(role=role.name, connection=connection, commit=False)
+                permission.grant(
+                    role=role.name, connection=connection, commit=False, feedback=feedback
+                )
         logger.info("All permissions granted to roles.")
         if commit:
+            if feedback:
+                feedback.lock_cancellation()
             connection.commit()

pum/schema_migrations.py

@@ -3,18 +3,28 @@ import logging
 import re
 
 import packaging
+import packaging.version
 import psycopg
 import psycopg.sql
 
-from .pum_config import PumConfig
-from .exceptions import PumException
+from .exceptions import PumSchemaMigrationError, PumSchemaMigrationNoBaselineError
 from .sql_content import SqlContent
+from .pum_config import PumConfig
 
 logger = logging.getLogger(__name__)
 
-MIGRATION_TABLE_VERSION = "2025.0"
+MIGRATION_TABLE_VERSION = 2  # Current schema version
 MIGRATION_TABLE_NAME = "pum_migrations"
 
+# TABLE VERSION HISTORY
+#
+# Version 1:
+#    Initial version with columns id, date_installed, module, version,
+#    beta_testing, changelog_files, parameters, migration_table_version
+#
+# Version 2:
+#   Changed migration_table_version type to integer and set module NOT NULL (version 2025.1 => 1, module 'tww')
+
 
 class SchemaMigrations:
     """Manage the schema migrations in the database.
@@ -36,6 +46,9 @@ class SchemaMigrations:
                 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.
@@ -61,8 +74,10 @@ class SchemaMigrations:
             "schema": psycopg.sql.Literal(self.config.config.pum.migration_table_schema),
         }
 
-        cursor = SqlContent(query).execute(connection, parameters=parameters)
-        return cursor.fetchone()[0]
+        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
 
     def exists_in_other_schemas(self, connection: psycopg.Connection) -> list[str]:
         """Check if the schema_migrations information table exists in other schemas.
@@ -85,8 +100,9 @@ class SchemaMigrations:
         parameters = {
             "schema": psycopg.sql.Literal(self.config.config.pum.migration_table_schema),
         }
-        cursor = SqlContent(query).execute(connection, parameters=parameters)
-        return [row[0] for row in cursor.fetchall()]
+        with connection.transaction():
+            cursor = SqlContent(query).execute(connection, parameters=parameters)
+            return [row[0] for row in (cursor._pum_results or [])]
 
     def create(
         self,
@@ -103,13 +119,13 @@ class SchemaMigrations:
                 distinct schemas. Default is false.
         """
         if self.exists(connection):
-            logger.info(
+            logger.debug(
                 f"{self.config.config.pum.migration_table_schema}.pum_migrations table already exists."
             )
             return
 
         if not allow_multiple_schemas and len(self.exists_in_other_schemas(connection)) > 0:
-            raise PumException(
+            raise PumSchemaMigrationError(
                 f"Another {self.config.config.pum.migration_table_schema}.{MIGRATION_TABLE_NAME} table exists in another schema (). "
                 "Please use the allow_multiple_schemas option to create a new one."
             )
@@ -130,30 +146,81 @@ class SchemaMigrations:
             (
             id uuid DEFAULT gen_random_uuid() PRIMARY KEY,
             date_installed timestamp without time zone NOT NULL DEFAULT now(),
-            module character varying(50), -- TODO: NOT NULL,
+            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 character varying(50) NOT NULL DEFAULT {version}
+            migration_table_version integer NOT NULL DEFAULT {version}
             );
         """
         )
 
-        comment_query = psycopg.sql.SQL(
-            "COMMENT ON TABLE {table} IS 'version: 1 --  schema_migration table 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 {parameters['schema']}.{parameters['table']} table")
+        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,
+        }
+
+        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,
@@ -175,24 +242,37 @@ class SchemaMigrations:
             commit: If true, the transaction is committed. The default is False.
 
         """
-        if isinstance(version, packaging.version.Version):
-            version = str(version)
+        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):
+        if not re.match(pattern, version_str):
             raise ValueError(f"Wrong version format: {version}. Must be x.y or x.y.z")
 
-        current = self.baseline(connection=connection)
-        if current and current >= version:
-            raise PumException(f"Cannot set baseline {version} as it is already set at {current}.")
+        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},
@@ -202,7 +282,8 @@ INSERT INTO {table} (
 
         query_parameters = {
             "table": self.migration_table_identifier,
-            "version": psycopg.sql.Literal(version),
+            "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 []),
@@ -210,11 +291,25 @@ INSERT INTO {table} (
         }
 
         logger.info(
-            f"Setting baseline version {version} in {self.config.config.pum.migration_table_schema}.{MIGRATION_TABLE_NAME} table"
+            f"Setting baseline version {version} in {self.migration_table_identifier_str} table"
         )
         SqlContent(code).execute(connection, parameters=query_parameters, commit=commit)
 
-    def baseline(self, connection: psycopg.Connection) -> str | None:
+    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.
+        """
+        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:
@@ -222,12 +317,17 @@ INSERT INTO {table} (
                 The database connection to get the baseline version.
 
         Returns:
-            str: The baseline version.
+            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
         """
 
         if not self.exists(connection=connection):
-            return None
+            raise PumSchemaMigrationError(
+                f"{self.migration_table_identifier_str} table does not exist."
+            )
 
         query = psycopg.sql.SQL(
             """
@@ -246,11 +346,14 @@ INSERT INTO {table} (
             "table": self.migration_table_identifier,
         }
 
-        cursor = SqlContent(query).execute(connection, parameters=parameters)
-        row = cursor.fetchone()
-        if row is None:
-            return None
-        return row[0]
+        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 PumSchemaMigrationNoBaselineError(
+                    f"Baseline version not found in the {self.migration_table_identifier_str} table."
+                )
+            return packaging.version.parse(row[0])
 
     def migration_details(self, connection: psycopg.Connection, version: str | None = None) -> dict:
         """Return the migration details from the migration table.
@@ -265,6 +368,8 @@ INSERT INTO {table} (
         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:
@@ -299,8 +404,40 @@ INSERT INTO {table} (
                 "version": psycopg.sql.Literal(version),
             }
 
-        cursor = SqlContent(query).execute(connection, parameters=parameters)
-        row = cursor.fetchone()
-        if row is None:
-            return None
-        return dict(zip([desc[0] for desc in cursor.description], row, strict=False))
+        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.
+        """
+
+        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