Core: Validate implementation classes loaded from properties

Author: akashchamp

pyiceberg/io/__init__.py

@@ -340,6 +340,8 @@ def _import_file_io(io_impl: str, properties: Properties) -> FileIO | None:
         module_name, class_name = ".".join(path_parts[:-1]), path_parts[-1]
         module = importlib.import_module(module_name)
         class_ = getattr(module, class_name)
+        if not isinstance(class_, type) or not issubclass(class_, FileIO):
+            raise ValueError(f"py-io-impl should be a subclass of FileIO, got: {io_impl}")
         return class_(properties)
     except ModuleNotFoundError:
         logger.warning(f"Could not initialize FileIO: {io_impl}", exc_info=logger.isEnabledFor(logging.DEBUG))

pyiceberg/io/pyarrow.py

@@ -232,6 +232,8 @@ def _import_retry_strategy(impl: str) -> S3RetryStrategy | None:
         module_name, class_name = ".".join(path_parts[:-1]), path_parts[-1]
         module = importlib.import_module(module_name)
         class_ = getattr(module, class_name)
+        if not isinstance(class_, type) or not issubclass(class_, S3RetryStrategy):
+            raise ValueError(f"retry-strategy-impl should be a subclass of S3RetryStrategy, got: {impl}")
         return class_()
     except (ModuleNotFoundError, AttributeError):
         warnings.warn(f"Could not initialize S3 retry strategy: {impl}", stacklevel=2)

pyiceberg/table/locations.py

@@ -178,6 +178,10 @@ def _import_location_provider(
         module_name, class_name = ".".join(path_parts[:-1]), path_parts[-1]
         module = importlib.import_module(module_name)
         class_ = getattr(module, class_name)
+        if not isinstance(class_, type) or not issubclass(class_, LocationProvider):
+            raise ValueError(
+                f"write.py-location-provider.impl should be a subclass of LocationProvider, got: {location_provider_impl}"
+            )
         return class_(table_location, table_properties)
     except ModuleNotFoundError:
         logger.warning(

tests/io/test_io.py

@@ -280,6 +280,11 @@ def test_import_file_io() -> None:
     assert isinstance(_import_file_io(ARROW_FILE_IO, {}), PyArrowFileIO)
 
 
+def test_import_file_io_wrong_type() -> None:
+    with pytest.raises(ValueError, match="py-io-impl should be a subclass of FileIO"):
+        _import_file_io("pyiceberg.table.locations.SimpleLocationProvider", {})
+
+
 def test_import_file_io_does_not_exist(caplog: Any) -> None:
     import logging
 

tests/io/test_pyarrow.py

@@ -3309,6 +3309,12 @@ def test_retry_strategy() -> None:
         io.new_input("s3://bucket/path/to/file")
 
 
+def test_retry_strategy_wrong_type() -> None:
+    io = PyArrowFileIO(properties={S3_RETRY_STRATEGY_IMPL: "pyiceberg.io.FileIO"})
+    with pytest.raises(ValueError, match="retry-strategy-impl should be a subclass of S3RetryStrategy"):
+        io.new_input("s3://bucket/path/to/file")
+
+
 def test_retry_strategy_not_found() -> None:
     io = PyArrowFileIO(properties={S3_RETRY_STRATEGY_IMPL: "pyiceberg.DoesNotExist"})
     with pytest.warns(UserWarning, match="Could not initialize S3 retry strategy: pyiceberg.DoesNotExist"):

tests/table/test_locations.py

@@ -60,6 +60,14 @@ def test_custom_location_provider() -> None:
     assert provider.new_data_location("my_file") == "custom_location_provider/my_file"
 
 
+def test_custom_location_provider_wrong_type() -> None:
+    with pytest.raises(ValueError, match="write.py-location-provider.impl should be a subclass of LocationProvider"):
+        load_location_provider(
+            table_location="table_location",
+            table_properties={"write.py-location-provider.impl": "pyiceberg.io.FileIO"},
+        )
+
+
 def test_custom_location_provider_single_path() -> None:
     with pytest.raises(ValueError, match=r"write\.py-location-provider\.impl should be full path"):
         load_location_provider(table_location="table_location", table_properties={"write.py-location-provider.impl": "not_found"})