From 216106c485d7cdbeb332d297a9a64dc79ee94300 Mon Sep 17 00:00:00 2001 From: Akshay Shah Date: Fri, 9 Oct 2026 12:58:17 -0500 Subject: [PATCH] fix(mlops): read disable_glue_table_creation from OfflineStoreConfig create_athena_query and construct_feature_group_to_be_merged looked up disable_glue_table_creation on DataCatalogConfig, which has no such field (it is on OfflineStoreConfig), so the getattr default False was always used. Feature groups created with DisableGlueTableCreation=True and a custom DataCatalogConfig.Catalog were therefore always queried against AwsDataCatalog instead of their catalog. Read the flag from offline_store_config. The unit tests now use real OfflineStoreConfig/DataCatalogConfig shapes and cover both values. --- .../mlops/feature_store/dataset_builder.py | 2 +- .../mlops/feature_store/feature_utils.py | 2 +- .../feature_store/test_dataset_builder.py | 22 +++++++++++++---- .../mlops/feature_store/test_feature_utils.py | 24 ++++++++++++++----- 4 files changed, 37 insertions(+), 13 deletions(-) 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 2805450998..f7e701f1cc 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/feature_store/dataset_builder.py +++ b/sagemaker-mlops/src/sagemaker/mlops/feature_store/dataset_builder.py @@ -161,7 +161,7 @@ def construct_feature_group_to_be_merged( raise RuntimeError(f"No metastore configured for FeatureGroup {fg.feature_group_name}.") catalog_config = fg.offline_store_config.data_catalog_config - disable_glue = getattr(catalog_config, "disable_glue_table_creation", False) or False + disable_glue = bool(fg.offline_store_config.disable_glue_table_creation) features = [fd.feature_name for fd in fg.feature_definitions] record_id = fg.record_identifier_feature_name diff --git a/sagemaker-mlops/src/sagemaker/mlops/feature_store/feature_utils.py b/sagemaker-mlops/src/sagemaker/mlops/feature_store/feature_utils.py index 4da8e8dc8e..1cca8cf997 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/feature_store/feature_utils.py +++ b/sagemaker-mlops/src/sagemaker/mlops/feature_store/feature_utils.py @@ -447,7 +447,7 @@ def create_athena_query(feature_group_name: str, session: Session): raise RuntimeError("No metastore is configured with this feature group.") catalog_config = fg.offline_store_config.data_catalog_config - disable_glue = getattr(catalog_config, "disable_glue_table_creation", False) or False + disable_glue = bool(fg.offline_store_config.disable_glue_table_creation) return AthenaQuery( catalog=catalog_config.catalog if disable_glue else "AwsDataCatalog", 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 e1361c3093..f7a813f6b0 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 @@ -7,6 +7,8 @@ from unittest.mock import Mock, patch, MagicMock import pandas as pd +from sagemaker.core.shapes import DataCatalogConfig, OfflineStoreConfig, S3StorageConfig + from sagemaker.mlops.feature_store import FeatureGroup from sagemaker.mlops.feature_store.dataset_builder import ( DatasetBuilder, @@ -102,8 +104,14 @@ def test_custom_join_settings(self): class TestConstructFeatureGroupToBeMerged: + @pytest.mark.parametrize( + "disable_glue_table_creation, expected_catalog", + [(False, "AwsDataCatalog"), (True, "MyCatalog")], + ) @patch("sagemaker.mlops.feature_store.dataset_builder.FeatureGroup") - def test_constructs_from_feature_group(self, mock_fg_class): + def test_constructs_from_feature_group( + self, mock_fg_class, disable_glue_table_creation, expected_catalog + ): mock_fg = MagicMock() mock_fg.feature_group_name = "test-fg" mock_fg.record_identifier_feature_name = "id" @@ -113,10 +121,13 @@ def test_constructs_from_feature_group(self, mock_fg_class): MagicMock(feature_name="value", feature_type="Fractional"), MagicMock(feature_name="event_time", feature_type="String"), ] - mock_fg.offline_store_config.data_catalog_config.catalog = "MyCatalog" - mock_fg.offline_store_config.data_catalog_config.database = "MyDatabase" - mock_fg.offline_store_config.data_catalog_config.table_name = "MyTable" - mock_fg.offline_store_config.data_catalog_config.disable_glue_table_creation = False + mock_fg.offline_store_config = OfflineStoreConfig( + s3_storage_config=S3StorageConfig(s3_uri="s3://bucket/offline"), + disable_glue_table_creation=disable_glue_table_creation, + data_catalog_config=DataCatalogConfig( + catalog="MyCatalog", database="MyDatabase", table_name="MyTable" + ), + ) mock_fg_class.get.return_value = mock_fg target_fg = MagicMock() @@ -127,6 +138,7 @@ def test_constructs_from_feature_group(self, mock_fg_class): included_feature_names=["id", "value"], ) + assert result.catalog == expected_catalog assert result.table_name == "MyTable" assert result.database == "MyDatabase" assert result.record_identifier_feature_name == "id" diff --git a/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_feature_utils.py b/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_feature_utils.py index fc18b76adc..caa564f96f 100644 --- a/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_feature_utils.py +++ b/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_feature_utils.py @@ -6,6 +6,8 @@ from unittest.mock import Mock, patch, MagicMock import pandas as pd +from sagemaker.core.shapes import DataCatalogConfig, OfflineStoreConfig, S3StorageConfig + from sagemaker.mlops.feature_store.feature_utils import ( load_feature_definitions_from_dataframe, as_hive_ddl, @@ -174,19 +176,29 @@ def test_custom_database_and_table(self, mock_fg_class): class TestCreateAthenaQuery: + @pytest.mark.parametrize( + "disable_glue_table_creation, expected_catalog", + [(False, "AwsDataCatalog"), (True, "MyCatalog")], + ) @patch("sagemaker.mlops.feature_store.feature_utils.CoreFeatureGroup") - def test_creates_athena_query(self, mock_fg_class): + def test_creates_athena_query( + self, mock_fg_class, disable_glue_table_creation, expected_catalog + ): mock_fg = MagicMock() - mock_fg.offline_store_config.data_catalog_config.catalog = "MyCatalog" - mock_fg.offline_store_config.data_catalog_config.database = "MyDatabase" - mock_fg.offline_store_config.data_catalog_config.table_name = "MyTable" - mock_fg.offline_store_config.data_catalog_config.disable_glue_table_creation = False + mock_fg.offline_store_config = OfflineStoreConfig( + s3_storage_config=S3StorageConfig(s3_uri="s3://bucket/offline"), + disable_glue_table_creation=disable_glue_table_creation, + data_catalog_config=DataCatalogConfig( + catalog="MyCatalog", database="MyDatabase", table_name="MyTable" + ), + ) mock_fg_class.get.return_value = mock_fg session = Mock() query = create_athena_query("my-fg", session) - assert query.catalog == "AwsDataCatalog" # disable_glue=False uses default + # The custom catalog is only used when Glue table creation is disabled. + assert query.catalog == expected_catalog assert query.database == "MyDatabase" assert query.table_name == "MyTable"