diff --git a/cumulusci/core/source_transforms/transforms.py b/cumulusci/core/source_transforms/transforms.py index 146802bb50..67e9ecb43a 100644 --- a/cumulusci/core/source_transforms/transforms.py +++ b/cumulusci/core/source_transforms/transforms.py @@ -2,6 +2,7 @@ import functools import io import os +import re import shutil import typing as T import zipfile @@ -338,10 +339,24 @@ def get_replace_string(self, context: TaskContext) -> str: ) try: - record_id = results["records"][0]["Id"] + # Extract the first column name using regular expression + column_name = re.search(r"\s*select\s+(\w+)", self.replace_record_id_query, re.IGNORECASE) + if column_name: + column_name = column_name.group(1) + else: + # Defaults to Id + column_name = "Id" + + except ValueError: + raise CumulusCIException( + "Unable to parse column name from replace_record_id_query. Please check format." + ) + + try: + record_id = results["records"][0][column_name] except KeyError: raise CumulusCIException( - "Results from the replace_record_id_query did not include an 'Id'. Please ensure the 'Id' field is included in your query's SELECT clause." + f"Results from the replace_record_id_query did not include an '{column_name}'. Please ensure the '{column_name}' field is included in your query's SELECT clause." ) return record_id