diff --git a/sagemaker-mlops/src/sagemaker/mlops/feature_store/dataset_builder.py b/sagemaker-mlops/src/sagemaker/mlops/feature_store/dataset_builder.py index 39a8bc9f5c..e7c6f9ef2d 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/feature_store/dataset_builder.py +++ b/sagemaker-mlops/src/sagemaker/mlops/feature_store/dataset_builder.py @@ -1,10 +1,13 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # Licensed under the Apache License, Version 2.0 """Dataset Builder for FeatureStore.""" + from dataclasses import dataclass, field from enum import Enum -from typing import Any, Dict, List, Union +from typing import Any, Dict, List, Optional, Union import datetime +import json +import logging import pandas as pd @@ -18,18 +21,26 @@ run_athena_query, ) +logger = logging.getLogger(__name__) + _DEFAULT_CATALOG = "AwsDataCatalog" _DEFAULT_DATABASE = "sagemaker_featurestore" _DTYPE_TO_FEATURE_TYPE = { - "object": "String", "string": "String", - "int64": "Integral", "int32": "Integral", - "float64": "Fractional", "float32": "Fractional", + "object": "String", + "string": "String", + "int64": "Integral", + "int32": "Integral", + "float64": "Fractional", + "float32": "Fractional", } _DTYPE_TO_ATHENA_TYPE = { - "object": "STRING", "int64": "INT", "float64": "DOUBLE", - "bool": "BOOLEAN", "datetime64[ns]": "TIMESTAMP", + "object": "STRING", + "int64": "INT", + "float64": "DOUBLE", + "bool": "BOOLEAN", + "datetime64[ns]": "TIMESTAMP", } @@ -90,6 +101,7 @@ class FeatureGroupToBeMerged: join_type (JoinTypeEnum): A JoinTypeEnum representing the type of join between the base and target feature groups. (default: JoinTypeEnum.INNER_JOIN). """ + features: List[str] included_feature_names: List[str] projected_feature_names: List[str] @@ -151,7 +163,7 @@ def construct_feature_group_to_be_merged( event_time_name = fg.event_time_feature_name event_time_type = next( (fd.feature_type for fd in fg.feature_definitions if fd.feature_name == event_time_name), - None + None, ) if feature_name_in_target and feature_name_in_target not in features: @@ -247,6 +259,7 @@ class DatasetBuilder: _event_time_identifier_feature_name: str = None _included_feature_names: List[str] = None _kms_key_id: str = None + _register_as_dataset: bool = False _event_time_identifier_feature_type: FeatureTypeEnum = None _point_in_time_accurate_join: bool = field(default=False, init=False) @@ -257,7 +270,10 @@ class DatasetBuilder: _write_time_ending_timestamp: datetime.datetime = field(default=None, init=False) _event_time_starting_timestamp: datetime.datetime = field(default=None, init=False) _event_time_ending_timestamp: datetime.datetime = field(default=None, init=False) - _feature_groups_to_be_merged: List[FeatureGroupToBeMerged] = field(default_factory=list, init=False) + _feature_groups_to_be_merged: List[FeatureGroupToBeMerged] = field( + default_factory=list, init=False + ) + _source_feature_groups: List[FeatureGroup] = field(default_factory=list, init=False) @classmethod def create( @@ -269,6 +285,7 @@ def create( event_time_identifier_feature_name: str = None, included_feature_names: List[str] = None, kms_key_id: str = None, + register_as_dataset: bool = False, ) -> "DatasetBuilder": """Create a DatasetBuilder for generating a Dataset. @@ -280,6 +297,12 @@ def create( event_time_identifier_feature_name: Required if base is DataFrame. included_feature_names: Features to include in output. kms_key_id: KMS key for encryption. + register_as_dataset: If True, register the output as a SM Dataset + (HubContent) after Athena completes, enabling automatic lineage + tracking between Feature Groups and Training Jobs. Requires + sagemaker:ImportHubContent permission. If the registration fails + due to missing permissions, a warning is logged and the CSV is + still returned. Default is False. Returns: DatasetBuilder instance. @@ -298,6 +321,7 @@ def create( _event_time_identifier_feature_name=event_time_identifier_feature_name, _included_feature_names=included_feature_names, _kms_key_id=kms_key_id, + _register_as_dataset=register_as_dataset, ) def with_feature_group( @@ -332,10 +356,15 @@ def with_feature_group( """ self._feature_groups_to_be_merged.append( construct_feature_group_to_be_merged( - feature_group, included_feature_names, target_feature_name_in_base, - feature_name_in_target, join_comparator, join_type, + feature_group, + included_feature_names, + target_feature_name_in_base, + feature_name_in_target, + join_comparator, + join_type, ) ) + self._source_feature_groups.append(feature_group) return self def point_in_time_accurate_join(self) -> "DatasetBuilder": @@ -431,7 +460,7 @@ def to_csv_file(self) -> tuple[str, str]: tuple: A tuple containing: - str: The S3 path of the .csv file - str: The query string executed - + Note: This method returns a tuple (csv_path, query_string). To get just the CSV path: csv_path, _ = builder.to_csv_file() @@ -450,7 +479,7 @@ def to_dataframe(self) -> tuple[pd.DataFrame, str]: tuple: A tuple containing: - pd.DataFrame: The pandas DataFrame object - str: The query string executed - + Note: This method returns a tuple (dataframe, query_string). To get just the DataFrame: df, _ = builder.to_dataframe() @@ -461,7 +490,6 @@ def to_dataframe(self) -> tuple[pd.DataFrame, str]: df = df.drop("row_recent", axis="columns") return df, query_string - def _to_csv_from_dataframe(self) -> tuple[str, str]: s3_folder, temp_table_name = upload_dataframe_to_s3( self._base, self._output_path, self._sagemaker_session, self._kms_key_id @@ -497,12 +525,22 @@ def _to_csv_from_dataframe(self) -> tuple[str, str]: def _to_csv_from_feature_group(self) -> tuple[str, str]: base_fg = construct_feature_group_to_be_merged(self._base, self._included_feature_names) self._record_identifier_feature_name = base_fg.record_identifier_feature_name - self._event_time_identifier_feature_name = base_fg.event_time_identifier_feature.feature_name - self._event_time_identifier_feature_type = base_fg.event_time_identifier_feature.feature_type + self._event_time_identifier_feature_name = ( + base_fg.event_time_identifier_feature.feature_name + ) + self._event_time_identifier_feature_type = ( + base_fg.event_time_identifier_feature.feature_type + ) query_string = self._construct_query_string(base_fg) result = self._run_query(query_string, base_fg.catalog, base_fg.database) - return self._extract_result(result) + csv_path, query = self._extract_result(result) + + if self._register_as_dataset: + query_execution_id = result.get("QueryExecution", {}).get("QueryExecutionId") + self._register_as_hub_content_dataset(csv_path, query_execution_id) + + return csv_path, query def _extract_result(self, query_result: dict) -> tuple[str, str]: execution = query_result.get("QueryExecution", {}) @@ -534,7 +572,6 @@ def _create_temp_table(self, temp_table_name: str, s3_folder: str): ) self._run_query(query, _DEFAULT_CATALOG, _DEFAULT_DATABASE) - def _construct_query_string(self, base: FeatureGroupToBeMerged) -> str: base_query = self._construct_table_query(base, "base") query = f"WITH fg_base AS ({base_query})" @@ -550,9 +587,7 @@ def _construct_query_string(self, base: FeatureGroupToBeMerged) -> str: selected += ", " + ", ".join( f'fg_{i}."{f}" as "{f}.{i+1}"' for f in fg.projected_feature_names ) - selected_final += ", " + ", ".join( - f'"{f}.{i+1}"' for f in fg.projected_feature_names - ) + selected_final += ", " + ", ".join(f'"{f}.{i+1}"' for f in fg.projected_feature_names) query += ( f"\nSELECT {selected_final}\nFROM (\n" @@ -594,7 +629,9 @@ def _construct_table_query(self, fg: FeatureGroupToBeMerged, suffix: str) -> str return ( f"SELECT {included}\n" f'FROM "{fg.database}"."{fg.table_name}" table_{suffix}\n' - + self._construct_where_query_string(suffix, fg.event_time_identifier_feature, ["NOT is_deleted"]) + + self._construct_where_query_string( + suffix, fg.event_time_identifier_feature, ["NOT is_deleted"] + ) ) if fg.table_type is TableType.FEATURE_GROUP and self._include_deleted_records: @@ -606,7 +643,9 @@ def _construct_table_query(self, fg: FeatureGroupToBeMerged, suffix: str) -> str f"{rank}) AS row_{suffix}\n" f'FROM "{fg.database}"."{fg.table_name}" origin_{suffix}\n' f"WHERE NOT is_deleted) AS table_{suffix}\n" - + self._construct_where_query_string(suffix, fg.event_time_identifier_feature, [f"row_{suffix} = 1"]) + + self._construct_where_query_string( + suffix, fg.event_time_identifier_feature, [f"row_{suffix} = 1"] + ) ) if fg.table_type is TableType.FEATURE_GROUP: @@ -626,7 +665,7 @@ def _construct_table_query(self, fg: FeatureGroupToBeMerged, suffix: str) -> str f"SELECT {included}\nFROM (\n" f"SELECT {included_with_write}\n" f'FROM "{fg.database}"."{fg.table_name}" table_{suffix}\n' - f"LEFT JOIN deleted_{suffix} ON table_{suffix}.\"{record_id}\" = deleted_{suffix}.\"{record_id}\"\n" + f'LEFT JOIN deleted_{suffix} ON table_{suffix}."{record_id}" = deleted_{suffix}."{record_id}"\n' f'WHERE deleted_{suffix}."{record_id}" IS NULL\n' f"UNION ALL\n" f"SELECT {included_with_write}\nFROM deleted_{suffix}\n" @@ -634,18 +673,20 @@ def _construct_table_query(self, fg: FeatureGroupToBeMerged, suffix: str) -> str f'ON table_{suffix}."{record_id}" = deleted_{suffix}."{record_id}"\n' f'AND (table_{suffix}."{event_time}" > deleted_{suffix}."{event_time}"\n{rank_cond})\n' f") AS table_{suffix}\n" - + self._construct_where_query_string(suffix, fg.event_time_identifier_feature, []) + + self._construct_where_query_string( + suffix, fg.event_time_identifier_feature, [] + ) ) return ( f"WITH {dedup},\n{deleted}\n" f"SELECT {included}\nFROM (\n" f"SELECT {included_with_write}\nFROM table_{suffix}\n" - f"LEFT JOIN deleted_{suffix} ON table_{suffix}.\"{record_id}\" = deleted_{suffix}.\"{record_id}\"\n" + f'LEFT JOIN deleted_{suffix} ON table_{suffix}."{record_id}" = deleted_{suffix}."{record_id}"\n' f'WHERE deleted_{suffix}."{record_id}" IS NULL\n' f"UNION ALL\n" f"SELECT {included_with_write}\nFROM deleted_{suffix}\n" - f"JOIN table_{suffix} ON table_{suffix}.\"{record_id}\" = deleted_{suffix}.\"{record_id}\"\n" + f'JOIN table_{suffix} ON table_{suffix}."{record_id}" = deleted_{suffix}."{record_id}"\n' f'AND (table_{suffix}."{event_time}" > deleted_{suffix}."{event_time}"\n{rank_cond})\n' f") AS table_{suffix}\n" + self._construct_where_query_string(suffix, fg.event_time_identifier_feature, []) @@ -672,7 +713,11 @@ def _construct_dedup_query(self, fg: FeatureGroupToBeMerged, suffix: str) -> str where_conds = [] if is_fg and self._write_time_ending_timestamp: where_conds.append(self._construct_write_time_condition(f"origin_{suffix}")) - where_conds.extend(self._construct_event_time_conditions(f"origin_{suffix}", fg.event_time_identifier_feature)) + where_conds.extend( + self._construct_event_time_conditions( + f"origin_{suffix}", fg.event_time_identifier_feature + ) + ) where_str = f"WHERE {' AND '.join(where_conds)}\n" if where_conds else "" dedup_where = f"WHERE dedup_row_{suffix} = 1\n" if is_fg else "" @@ -693,7 +738,9 @@ def _construct_deleted_query(self, fg: FeatureGroupToBeMerged, suffix: str) -> s rank = f'ORDER BY origin_{suffix}."{event_time}" DESC' if fg.table_type is TableType.FEATURE_GROUP: - rank += f', origin_{suffix}."api_invocation_time" DESC, origin_{suffix}."write_time" DESC\n' + rank += ( + f', origin_{suffix}."api_invocation_time" DESC, origin_{suffix}."write_time" DESC\n' + ) write_cond = "" if fg.table_type is TableType.FEATURE_GROUP and self._write_time_ending_timestamp: @@ -701,7 +748,9 @@ def _construct_deleted_query(self, fg: FeatureGroupToBeMerged, suffix: str) -> s event_conds = "" if self._event_time_starting_timestamp and self._event_time_ending_timestamp: - conds = self._construct_event_time_conditions(f"origin_{suffix}", fg.event_time_identifier_feature) + conds = self._construct_event_time_conditions( + f"origin_{suffix}", fg.event_time_identifier_feature + ) event_conds = "".join(f"AND {c}\n" for c in conds) return ( @@ -723,7 +772,9 @@ def _construct_where_query_string( if isinstance(self._base, FeatureGroup) and self._write_time_ending_timestamp: conditions.append(self._construct_write_time_condition(f"table_{suffix}")) - conditions.extend(self._construct_event_time_conditions(f"table_{suffix}", event_time_feature)) + conditions.extend( + self._construct_event_time_conditions(f"table_{suffix}", event_time_feature) + ) return f"WHERE {' AND '.join(conditions)}" if conditions else "" def _validate_options(self): @@ -736,16 +787,26 @@ def _validate_options(self): raise ValueError("number_of_records must be non-negative.") if is_df_base and no_joins: if self._include_deleted_records: - raise ValueError("include_deleted_records() only works for FeatureGroup if no join.") + raise ValueError( + "include_deleted_records() only works for FeatureGroup if no join." + ) if self._include_duplicated_records: - raise ValueError("include_duplicated_records() only works for FeatureGroup if no join.") + raise ValueError( + "include_duplicated_records() only works for FeatureGroup if no join." + ) if self._write_time_ending_timestamp: raise ValueError("as_of() only works for FeatureGroup if no join.") if self._point_in_time_accurate_join and no_joins: raise ValueError("point_in_time_accurate_join() requires at least one join.") - def _construct_event_time_conditions(self, table: str, event_time_feature: FeatureDefinition) -> List[str]: - cast_fn = "from_iso8601_timestamp" if event_time_feature.feature_type == FeatureTypeEnum.STRING else "from_unixtime" + def _construct_event_time_conditions( + self, table: str, event_time_feature: FeatureDefinition + ) -> List[str]: + cast_fn = ( + "from_iso8601_timestamp" + if event_time_feature.feature_type == FeatureTypeEnum.STRING + else "from_unixtime" + ) conditions = [] if self._event_time_starting_timestamp: conditions.append( @@ -761,7 +822,7 @@ def _construct_event_time_conditions(self, table: str, event_time_feature: Featu def _construct_write_time_condition(self, table: str) -> str: ts = self._write_time_ending_timestamp.replace(microsecond=0) - return f'{table}."write_time" <= to_timestamp(\'{ts}\', \'yyyy-mm-dd hh24:mi:ss\')' + return f"{table}.\"write_time\" <= to_timestamp('{ts}', 'yyyy-mm-dd hh24:mi:ss')" def _construct_join_condition(self, fg: FeatureGroupToBeMerged, suffix: str) -> str: target_feature = fg.feature_name_in_target or fg.record_identifier_feature_name @@ -771,11 +832,123 @@ def _construct_join_condition(self, fg: FeatureGroupToBeMerged, suffix: str) -> ) if self._point_in_time_accurate_join: - base_cast = "from_iso8601_timestamp" if self._event_time_identifier_feature_type == FeatureTypeEnum.STRING else "from_unixtime" - fg_cast = "from_iso8601_timestamp" if fg.event_time_identifier_feature.feature_type == FeatureTypeEnum.STRING else "from_unixtime" + base_cast = ( + "from_iso8601_timestamp" + if self._event_time_identifier_feature_type == FeatureTypeEnum.STRING + else "from_unixtime" + ) + fg_cast = ( + "from_iso8601_timestamp" + if fg.event_time_identifier_feature.feature_type == FeatureTypeEnum.STRING + else "from_unixtime" + ) join += ( f'\nAND {base_cast}(fg_base."{self._event_time_identifier_feature_name}") >= ' f'{fg_cast}(fg_{suffix}."{fg.event_time_identifier_feature.feature_name}")' ) return join + + def _collect_source_feature_group_arns(self) -> List[str]: + """Collect ARNs of all source Feature Groups used in this dataset extraction. + + Returns: + List of Feature Group ARNs (base + any merged FGs). + """ + arns = [] + + # Base Feature Group + if isinstance(self._base, FeatureGroup): + fg_arn = getattr(self._base, "feature_group_arn", None) + if fg_arn: + arns.append(str(fg_arn)) + + # Merged Feature Groups (tracked via with_feature_group()) + for fg in self._source_feature_groups: + fg_arn = getattr(fg, "feature_group_arn", None) + if fg_arn and str(fg_arn) not in arns: + arns.append(str(fg_arn)) + + return arns + + def _register_as_hub_content_dataset( + self, csv_path: str, query_execution_id: Optional[str] = None + ) -> None: + """Register the output CSV as a SM Dataset (HubContent) for lineage tracking. + + Creates a HubContent of type DataSet with source Feature Group ARNs in the + metadata. When Eureka processes this HubContent creation event, it automatically + creates FG → Dataset lineage associations. + + This is a best-effort operation: if it fails due to missing permissions + (AccessDeniedException), a warning is logged and the method returns without + raising — the primary workflow (returning the CSV) is not affected. + + Args: + csv_path: S3 path of the generated CSV file. + query_execution_id: Athena query execution ID (for provenance tracking). + """ + source_fg_arns = self._collect_source_feature_group_arns() + if not source_fg_arns: + logger.warning( + "register_as_dataset=True but no Feature Group ARNs found. " + "Skipping dataset registration." + ) + return + + # Build the HubContent document with source FG metadata + hub_content_document = json.dumps( + { + "DatasetS3Uri": csv_path, + "ContentMetadata": { + "sourceFeatureGroups": source_fg_arns, + "extractionMethod": "FeatureStoreDatasetBuilder", + "athenaQueryExecutionId": query_execution_id or "", + }, + } + ) + + # Generate a dataset name from the base FG name + timestamp + base_name = "" + if isinstance(self._base, FeatureGroup): + base_name = getattr(self._base, "feature_group_name", "dataset") + else: + base_name = "dataframe-dataset" + timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S") + dataset_name = f"fs-{base_name}-{timestamp}" + + try: + from sagemaker.core.resources import HubContent + + HubContent.import_hub_content( + hub_content_name=dataset_name, + hub_content_version="1", + hub_content_type="DataSet", + document_schema_version="4.0.0", + hub_name="AIRegistry", + hub_content_document=hub_content_document, + hub_content_description=f"Dataset extracted from Feature Groups: {', '.join(source_fg_arns)}", + session=self._sagemaker_session, + ) + logger.info( + "Registered dataset '%s' as SM Dataset (HubContent) with source FGs: %s", + dataset_name, + source_fg_arns, + ) + except Exception as e: + # Graceful fallback: log warning, don't block the primary workflow + error_msg = str(e) + if "AccessDenied" in error_msg or "not authorized" in error_msg.lower(): + logger.warning( + "Unable to register dataset as HubContent due to missing permissions " + "(sagemaker:ImportHubContent). Lineage will not be created for this " + "dataset extraction. To enable lineage, add sagemaker:ImportHubContent " + "permission to your execution role. Error: %s", + error_msg, + ) + else: + logger.warning( + "Failed to register dataset as HubContent. Lineage will not be created. " + "Error: %s", + error_msg, + ) diff --git a/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_dataset_builder.py b/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_dataset_builder.py index 4cd2e47e08..ad89d28333 100644 --- a/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_dataset_builder.py +++ b/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_dataset_builder.py @@ -1,6 +1,7 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # Licensed under the Apache License, Version 2.0 """Unit tests for dataset_builder.py""" + import datetime import pytest from unittest.mock import Mock, patch, MagicMock @@ -17,7 +18,6 @@ ) from sagemaker.mlops.feature_store.feature_definition import ( FeatureDefinition, - FeatureTypeEnum, ) @@ -152,11 +152,13 @@ def mock_session(self): @pytest.fixture def sample_dataframe(self): - return pd.DataFrame({ - "id": [1, 2, 3], - "value": [1.1, 2.2, 3.3], - "event_time": ["2024-01-01", "2024-01-02", "2024-01-03"], - }) + return pd.DataFrame( + { + "id": [1, 2, 3], + "value": [1.1, 2.2, 3.3], + "event_time": ["2024-01-01", "2024-01-02", "2024-01-03"], + } + ) def test_initialization_with_dataframe(self, mock_session, sample_dataframe): builder = DatasetBuilder( @@ -345,3 +347,176 @@ def test_to_csv_raises_for_invalid_base(self, mock_session): with pytest.raises(ValueError, match="must be either"): builder.to_csv_file() + + +class TestDatasetBuilderRegisterAsDataset: + """Tests for the register_as_dataset lineage feature.""" + + @pytest.fixture + def mock_session(self): + return Mock() + + @pytest.fixture + def mock_feature_group(self): + fg = MagicMock(spec=FeatureGroup) + fg.feature_group_name = "customers-fg" + fg.feature_group_arn = "arn:aws:sagemaker:us-west-2:123456789012:feature-group/customers-fg" + return fg + + def test_register_as_dataset_default_false(self, mock_session): + """Default register_as_dataset is False.""" + builder = DatasetBuilder.create( + base=MagicMock(spec=FeatureGroup), + output_path="s3://bucket/output", + session=mock_session, + ) + assert builder._register_as_dataset is False + + def test_register_as_dataset_true_sets_flag(self, mock_session): + """register_as_dataset=True is stored correctly.""" + builder = DatasetBuilder.create( + base=MagicMock(spec=FeatureGroup), + output_path="s3://bucket/output", + session=mock_session, + register_as_dataset=True, + ) + assert builder._register_as_dataset is True + + def test_collect_source_fg_arns_from_base(self, mock_session, mock_feature_group): + """Collects base FG ARN.""" + builder = DatasetBuilder( + _sagemaker_session=mock_session, + _base=mock_feature_group, + _output_path="s3://bucket/output", + _register_as_dataset=True, + ) + arns = builder._collect_source_feature_group_arns() + assert arns == ["arn:aws:sagemaker:us-west-2:123456789012:feature-group/customers-fg"] + + def test_collect_source_fg_arns_with_merged_fg(self, mock_session, mock_feature_group): + """Collects base + merged FG ARNs.""" + merged_fg = MagicMock(spec=FeatureGroup) + merged_fg.feature_group_arn = ( + "arn:aws:sagemaker:us-west-2:123456789012:feature-group/orders-fg" + ) + + builder = DatasetBuilder( + _sagemaker_session=mock_session, + _base=mock_feature_group, + _output_path="s3://bucket/output", + _register_as_dataset=True, + ) + builder._source_feature_groups.append(merged_fg) + + arns = builder._collect_source_feature_group_arns() + assert len(arns) == 2 + assert "arn:aws:sagemaker:us-west-2:123456789012:feature-group/customers-fg" in arns + assert "arn:aws:sagemaker:us-west-2:123456789012:feature-group/orders-fg" in arns + + def test_collect_source_fg_arns_deduplicates(self, mock_session, mock_feature_group): + """Doesn't duplicate ARNs if same FG used twice.""" + builder = DatasetBuilder( + _sagemaker_session=mock_session, + _base=mock_feature_group, + _output_path="s3://bucket/output", + _register_as_dataset=True, + ) + # Add same FG as merged + builder._source_feature_groups.append(mock_feature_group) + + arns = builder._collect_source_feature_group_arns() + assert len(arns) == 1 + + def test_collect_source_fg_arns_dataframe_base_empty(self, mock_session): + """DataFrame base has no FG ARN.""" + df = pd.DataFrame({"id": [1], "event_time": ["2024-01-01"]}) + builder = DatasetBuilder( + _sagemaker_session=mock_session, + _base=df, + _output_path="s3://bucket/output", + _record_identifier_feature_name="id", + _event_time_identifier_feature_name="event_time", + _register_as_dataset=True, + ) + arns = builder._collect_source_feature_group_arns() + assert arns == [] + + def test_register_hub_content_called_on_success(self, mock_session, mock_feature_group): + """HubContent.import_hub_content is called with correct params.""" + builder = DatasetBuilder( + _sagemaker_session=mock_session, + _base=mock_feature_group, + _output_path="s3://bucket/output", + _register_as_dataset=True, + ) + + with patch( + "sagemaker.core.resources.HubContent", + ) as mock_hc: + builder._register_as_hub_content_dataset( + csv_path="s3://bucket/output/result.csv", + query_execution_id="abc-123", + ) + mock_hc.import_hub_content.assert_called_once() + call_kwargs = mock_hc.import_hub_content.call_args[1] + assert call_kwargs["hub_content_type"] == "DataSet" + assert call_kwargs["hub_name"] == "AIRegistry" + assert "customers-fg" in call_kwargs["hub_content_name"] + assert mock_session == call_kwargs["session"] + + def test_register_graceful_on_access_denied(self, mock_session, mock_feature_group, caplog): + """AccessDeniedException logs warning, doesn't raise.""" + builder = DatasetBuilder( + _sagemaker_session=mock_session, + _base=mock_feature_group, + _output_path="s3://bucket/output", + _register_as_dataset=True, + ) + + with patch( + "sagemaker.core.resources.HubContent", + ) as mock_hc: + mock_hc.import_hub_content.side_effect = Exception("AccessDenied: not authorized") + # Should NOT raise + builder._register_as_hub_content_dataset( + csv_path="s3://bucket/output/result.csv", + ) + + def test_register_graceful_on_generic_error(self, mock_session, mock_feature_group): + """Generic errors log warning, don't raise.""" + builder = DatasetBuilder( + _sagemaker_session=mock_session, + _base=mock_feature_group, + _output_path="s3://bucket/output", + _register_as_dataset=True, + ) + + with patch( + "sagemaker.core.resources.HubContent", + ) as mock_hc: + mock_hc.import_hub_content.side_effect = Exception("Some service error") + # Should NOT raise + builder._register_as_hub_content_dataset( + csv_path="s3://bucket/output/result.csv", + ) + + def test_register_skipped_when_no_fg_arns(self, mock_session): + """Skips registration when no FG ARNs available (DataFrame base, no merged FGs).""" + df = pd.DataFrame({"id": [1], "event_time": ["2024-01-01"]}) + builder = DatasetBuilder( + _sagemaker_session=mock_session, + _base=df, + _output_path="s3://bucket/output", + _record_identifier_feature_name="id", + _event_time_identifier_feature_name="event_time", + _register_as_dataset=True, + ) + + with patch( + "sagemaker.core.resources.HubContent", + ) as mock_hc: + builder._register_as_hub_content_dataset( + csv_path="s3://bucket/output/result.csv", + ) + # Should NOT call import_hub_content since no FG ARNs + mock_hc.import_hub_content.assert_not_called()