pum 1.0.0__py3-none-any.whl

pum/role_manager.py

@@ -0,0 +1,253 @@
+import enum
+from typing import Optional
+import copy
+import psycopg
+import logging
+
+from .sql_content import SqlContent
+
+logger = logging.getLogger(__name__)
+
+
+class PermissionType(enum.Enum):
+    """Enum for permission types.
+
+    Attributes:
+        READ (str): Read permission.
+        WRITE (str): Write permission.
+    """
+
+    READ = "read"
+    WRITE = "write"
+
+
+class Permission:
+    """Class to represent a permission for a database role.
+
+    Attributes:
+        type: Type of permission (read or write).
+        schemas: List of schemas this permission applies to.
+    """
+
+    def __init__(self, type: PermissionType | str, schemas: list[str] = None) -> None:
+        if not isinstance(type, PermissionType):
+            type = PermissionType(type)
+        self.type = type
+        self.schemas = schemas
+
+    def grant(
+        self,
+        role: str,
+        connection: psycopg.Connection,
+        commit: bool = False,
+    ) -> 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.
+        """
+        if not isinstance(role, str):
+            raise TypeError("Role must be a string.")
+
+        if not self.schemas:
+            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 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};
+                        ALTER DEFAULT PRIVILEGES IN SCHEMA {schema} GRANT SELECT, REFERENCES, TRIGGER ON TABLES TO {role};
+                           """).execute(
+                    connection=connection,
+                    commit=False,
+                    parameters={
+                        "schema": psycopg.sql.Identifier(schema),
+                        "role": psycopg.sql.Identifier(role),
+                    },
+                )
+            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};
+                        ALTER DEFAULT PRIVILEGES IN SCHEMA {schema} GRANT ALL ON TABLES TO {role};
+                           """).execute(
+                    connection=connection,
+                    commit=False,
+                    parameters={
+                        "schema": psycopg.sql.Identifier(schema),
+                        "role": psycopg.sql.Identifier(role),
+                    },
+                )
+            else:
+                raise ValueError(f"Unknown permission type: {self.type}")
+
+        if commit:
+            connection.commit()
+
+
+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.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        permissions: list[Permission] | list[str],
+        *,
+        inherit: Optional["Role"] = None,
+        description: str | None = None,
+    ) -> None:
+        """Initialize the Role class.
+        Args:
+            name: Name of the role.
+            permissions: List of permissions associated with the role.
+            inherit: Optional role to inherit permissions from.
+            description: Optional description of the role.
+        """
+        self.name = name
+        if isinstance(permissions, list) and all(isinstance(p, dict) for p in permissions):
+            self._permissions = [Permission(**p) for p in permissions]
+        elif isinstance(permissions, list) and all(isinstance(p, Permission) for p in permissions):
+            self._permissions = permissions
+        else:
+            raise TypeError("Permissions must be a list of dictionnaries or Permission instances.")
+
+        if inherit is not None and not isinstance(inherit, Role):
+            raise TypeError("Inherit must be a Role instance or None.")
+        self.inherit = inherit
+        self.description = description
+
+    def permissions(self):
+        """
+        Returns the list of permissions associated with the role.
+        """
+        return self._permissions
+
+    def exists(self, connection: psycopg.Connection) -> bool:
+        """Check if the role exists in the database.
+        Args:
+            connection: The database connection to execute the SQL statements.
+        Returns:
+            bool: True if the role exists, False otherwise.
+        """
+        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
+
+    def create(
+        self, connection: psycopg.Connection, grant: bool = False, commit: bool = False
+    ) -> 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.
+        """
+        if self.exists(connection):
+            logger.info(f"Role {self.name} already exists, skipping creation.")
+        else:
+            logger.info(f"Creating role {self.name}.")
+            SqlContent(
+                "CREATE ROLE {name} NOSUPERUSER INHERIT NOCREATEDB NOCREATEROLE NOREPLICATION"
+            ).execute(
+                connection=connection,
+                commit=False,
+                parameters={"name": psycopg.sql.Identifier(self.name)},
+            )
+            if self.description:
+                SqlContent("COMMENT ON ROLE {name} IS {description}").execute(
+                    connection=connection,
+                    commit=False,
+                    parameters={
+                        "name": psycopg.sql.Identifier(self.name),
+                        "description": psycopg.sql.Literal(self.description),
+                    },
+                )
+            if self.inherit:
+                SqlContent("GRANT {inherit} TO {role}").execute(
+                    connection=connection,
+                    commit=False,
+                    parameters={
+                        "inherit": psycopg.sql.Identifier(self.inherit.name),
+                        "role": psycopg.sql.Identifier(self.name),
+                    },
+                )
+        if grant:
+            for permission in self.permissions():
+                permission.grant(role=self.name, connection=connection, commit=commit)
+
+        if commit:
+            connection.commit()
+
+
+class RoleManager:
+    """
+    RoleManager manages a collection of Role objects,
+    allowing creation and permission management
+    for multiple roles in the PostgreSQL database.
+    """
+
+    def __init__(self, roles=list[Role] | list[dict]) -> None:
+        """Initialize the RoleManager class.:
+        Args:
+            roles: List of roles or dictionaries defining roles.
+            Each role can be a dictionary with keys 'name', 'permissions', and optional 'description' and 'inherit'.
+        """
+        if isinstance(roles, list) and all(isinstance(role, dict) for role in roles):
+            self.roles = {}
+            for role in roles:
+                _inherit = role.get("inherit")
+                if _inherit is not None:
+                    if _inherit not in self.roles:
+                        raise ValueError(
+                            f"Inherited role {_inherit} does not exist in the already defined roles. Pay attention to the order of the roles in the list."
+                        )
+                    role["inherit"] = self.roles[_inherit]
+                self.roles[role["name"]] = Role(**role)
+        elif isinstance(roles, list) and all(isinstance(role, Role) for role in roles):
+            _roles = copy.deepcopy(roles)
+            self.roles = {role.name: role for role in _roles}
+        else:
+            raise TypeError("Roles must be a list of dictionaries or Role instances.")
+
+        for role in self.roles.values():
+            if role.inherit is not None and role.inherit not in self.roles.values():
+                raise ValueError(
+                    f"Inherited role {role.inherit.name} does not exist in the defined roles."
+                )
+
+    def create_roles(
+        self, connection: psycopg.Connection, grant: bool = False, commit: bool = False
+    ) -> 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.
+        """
+        for role in self.roles.values():
+            role.create(connection=connection, commit=False, grant=grant)
+        if commit:
+            connection.commit()
+
+    def grant_permissions(self, connection: psycopg.Connection, commit: bool = False) -> 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.
+        """
+        for role in self.roles.values():
+            for permission in role.permissions():
+                permission.grant(role=role.name, connection=connection, commit=False)
+        if commit:
+            connection.commit()

pum/schema_migrations.py

@@ -0,0 +1,306 @@
+import json
+import logging
+import re
+
+import packaging
+import psycopg
+import psycopg.sql
+
+from .pum_config import PumConfig
+from .exceptions import PumException
+from .sql_content import SqlContent
+
+logger = logging.getLogger(__name__)
+
+MIGRATION_TABLE_VERSION = "2025.0"
+MIGRATION_TABLE_NAME = "pum_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.
+    """
+
+    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),
+            ]
+        )
+
+    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),
+        }
+
+        cursor = SqlContent(query).execute(connection, parameters=parameters)
+        return cursor.fetchone()[0]
+
+    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.
+
+        """
+        query = psycopg.sql.SQL(
+            """
+            SELECT table_schema
+            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),
+        }
+        cursor = SqlContent(query).execute(connection, parameters=parameters)
+        return [row[0] for row in cursor.fetchall()]
+
+    def create(
+        self,
+        connection: psycopg.Connection,
+        *,
+        allow_multiple_schemas: 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_schemas: If true, several pum_migrations tables are allowed in
+                distinct schemas. Default is false.
+        """
+        if self.exists(connection):
+            logger.info(
+                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(
+                f"Another {self.config.config.pum.migration_table_schema}.{self.config.config.pum.migration_table_name} table exists in another schema (). "
+                "Please use the allow_multiple_schemas option to create a new one."
+            )
+
+        # 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), -- TODO: 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}
+            );
+        """
+        )
+
+        comment_query = psycopg.sql.SQL(
+            "COMMENT ON TABLE {table} IS 'version: 1 --  schema_migration table version';"
+        )
+
+        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")
+
+        if commit:
+            connection.commit()
+
+    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.
+
+        """
+        if isinstance(version, packaging.version.Version):
+            version = str(version)
+        pattern = re.compile(r"^\d+\.\d+\.\d+$")
+        if not re.match(pattern, version):
+            raise ValueError(f"Wrong version format: {version}. Must be x.x.x")
+
+        current = self.baseline(connection=connection)
+        if current and current >= version:
+            raise PumException(f"Cannot set baseline {version} as it is already set at {current}.")
+
+        code = psycopg.sql.SQL("""
+INSERT INTO {table} (
+    version,
+    beta_testing,
+    migration_table_version,
+    changelog_files,
+    parameters
+) VALUES (
+    {version},
+    {beta_testing},
+    {migration_table_version},
+    {changelog_files},
+    {parameters}
+);""")
+
+        query_parameters = {
+            "table": self.migration_table_identifier,
+            "version": psycopg.sql.Literal(version),
+            "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.config.config.pum.migration_table_schema}.{MIGRATION_TABLE_NAME} table"
+        )
+        SqlContent(code).execute(connection, parameters=query_parameters, commit=commit)
+
+    def baseline(self, connection: psycopg.Connection) -> str | None:
+        """Return the baseline version from the migration table.
+
+        Args:
+            connection: psycopg.Connection
+                The database connection to get the baseline version.
+
+        Returns:
+            str: The baseline version.
+
+        """
+
+        if not self.exists(connection=connection):
+            return None
+
+        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.fetchone()
+        if row is None:
+            return None
+        return row[0]
+
+    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.
+
+        """
+        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),
+            }
+
+        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))