Merge pull request #45 from datamasque/DM-4094-tagging-role

feat: DM-4094: Add iam_role_arn to DatabaseConnectionConfig

Author: alxboyle

HISTORY.rst

@@ -2,6 +2,15 @@
 History
 =======
 
+1.2.5 (2026-08-17)
+------------------
+
+* Added ``iam_role_arn`` to ``DatabaseConnectionConfig``: the IAM role DataMasque assumes to tag
+  the connection's AWS resource, for resources in another AWS account. Sent only for the engines
+  that can be an Amazon RDS instance, Aurora cluster or Redshift cluster.
+
+Requires server version 3.26.16
+
 1.2.4 (2026-08-12)
 ------------------
 

datamasque/client/models/connection.py

@@ -289,6 +289,22 @@ def _strip_encrypted_password(cls, data: dict) -> dict:
         return data
 
 
+# Engines that can sit behind an Amazon RDS, Aurora or Redshift endpoint, and so can carry a role
+# for DataMasque to assume when tagging the resource. Aurora needs no entry of its own: it uses
+# these same engines with a cluster endpoint.
+_AWS_TAGGABLE_DATABASE_TYPES = frozenset(
+    {
+        DatabaseType.postgres,
+        DatabaseType.mysql,
+        DatabaseType.mariadb,
+        DatabaseType.oracle,
+        DatabaseType.mssql,
+        DatabaseType.db2_luw,
+        DatabaseType.redshift,
+    }
+)
+
+
 class DatabaseConnectionConfig(ConnectionConfig):
     """
     Connection configuration for a SQL database.
@@ -309,6 +325,11 @@ class DatabaseConnectionConfig(ConnectionConfig):
     is_read_only: bool = False
     s3_bucket_name: Optional[str] = None
     s3_redshift_iam_role: Optional[str] = None
+    # An IAM role for DataMasque to assume when tagging this database's AWS resource after a run.
+    # Only the engines in `_AWS_TAGGABLE_DATABASE_TYPES` can be an RDS instance, Aurora cluster or
+    # Redshift cluster, so it is pruned for the rest. Distinct from `s3_redshift_iam_role`, which is
+    # the role the Redshift cluster itself uses to reach S3.
+    iam_role_arn: Optional[str] = None
 
     @model_validator(mode="after")
     def _reject_special_engines(self) -> "DatabaseConnectionConfig":
@@ -348,6 +369,8 @@ def _serialize(self, handler: Callable) -> dict:
         if db_type is not DatabaseType.redshift:
             d.pop("s3_bucket_name", None)
             d.pop("s3_redshift_iam_role", None)
+        if db_type not in _AWS_TAGGABLE_DATABASE_TYPES:
+            d.pop("iam_role_arn", None)
         if not d.get("engine_options"):
             d.pop("engine_options", None)
         return d

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "datamasque-python"
-version = "1.2.4"
+version = "1.2.5"
 description = "Official Python client for the DataMasque data-masking API."
 authors = [
     { name = "DataMasque Ltd" },

tests/test_connections.py

@@ -1400,3 +1400,86 @@ def test_create_or_update_mongo_connection(client):
     assert sent["db_type"] == "mongodb"
     assert sent["dbpassword"] == "hunter2"
     assert sent["replica_set"] == "rs0"
+
+
+def _database_connection(database_type: DatabaseType, **extra) -> DatabaseConnectionConfig:
+    return DatabaseConnectionConfig(
+        name="conn",
+        host="mydb.abc123.ap-southeast-2.rds.amazonaws.com",
+        port=5432,
+        database="db",
+        user="user",
+        password="secret",
+        database_type=database_type,
+        **extra,
+    )
+
+
+@pytest.mark.parametrize(
+    "database_type",
+    [
+        DatabaseType.postgres,
+        DatabaseType.mysql,
+        DatabaseType.mariadb,
+        DatabaseType.oracle,
+        DatabaseType.mssql,
+        DatabaseType.db2_luw,
+        DatabaseType.redshift,
+    ],
+)
+def test_database_connection_sends_tagging_iam_role_for_aws_hosted_engines(database_type):
+    """The engines that can sit behind an RDS, Aurora or Redshift endpoint can carry a tagging role."""
+    extra = (
+        {"s3_bucket_name": "bucket", "s3_redshift_iam_role": "arn:aws:iam::123456789012:role/redshift-s3"}
+        if database_type is DatabaseType.redshift
+        else {}
+    )
+    conn = _database_connection(database_type, iam_role_arn="arn:aws:iam::119836602066:role/tagger", **extra)
+
+    api_dict = conn.model_dump(exclude_none=True, by_alias=True, mode="json")
+
+    assert api_dict["iam_role_arn"] == "arn:aws:iam::119836602066:role/tagger"
+
+
+@pytest.mark.parametrize(
+    "database_type",
+    [
+        DatabaseType.db2i,
+        DatabaseType.informix,
+        DatabaseType.saphana,
+        # The two worth documenting rather than merely covering. `mssql_linked` inherits the field
+        # through `MssqlLinkedServerConnectionConfig` and is excluded only by its own enum value;
+        # `databricks_lakebase` is AWS-hosted, so it reads like an omission from the taggable set
+        # until you know it presents no endpoint tagging can act on.
+        DatabaseType.mssql_linked,
+        DatabaseType.databricks_lakebase,
+    ],
+)
+def test_database_connection_omits_tagging_iam_role_for_engines_that_cannot_be_aws_resources(database_type):
+    """Mirrors how `s3_redshift_iam_role` is pruned: the server has no such field on these engines."""
+    conn = _database_connection(database_type, iam_role_arn="arn:aws:iam::119836602066:role/tagger")
+
+    api_dict = conn.model_dump(exclude_none=True, by_alias=True, mode="json")
+
+    assert "iam_role_arn" not in api_dict
+
+
+def test_database_connection_omits_tagging_iam_role_when_unset():
+    conn = _database_connection(DatabaseType.postgres)
+
+    assert "iam_role_arn" not in conn.model_dump(exclude_none=True, by_alias=True, mode="json")
+
+
+def test_redshift_connection_keeps_its_cluster_s3_role_separate_from_its_tagging_role():
+    """Two ARNs with different meanings: the cluster's own S3 access, and the role DataMasque assumes to tag."""
+    conn = _database_connection(
+        DatabaseType.redshift,
+        s3_bucket_name="bucket",
+        s3_redshift_iam_role="arn:aws:iam::123456789012:role/redshift-s3",
+        iam_role_arn="arn:aws:iam::119836602066:role/tagger",
+    )
+
+    api_dict = conn.model_dump(exclude_none=True, by_alias=True, mode="json")
+
+    assert api_dict["s3_redshift_iam_role"] == "arn:aws:iam::123456789012:role/redshift-s3"
+    assert api_dict["iam_role_arn"] == "arn:aws:iam::119836602066:role/tagger"