fix(mlops): read disable_glue_table_creation from OfflineStoreConfig

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

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",

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"

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"