SavedDatasetStorage registry: custom_storage is claimed by three classes, so from_proto dispatch depends on import order #6946

Problem

_StorageRegistry keys storage classes by _proto_attr_name in a plain dict:

# sdk/python/feast/saved_dataset.py
class _StorageRegistry(type):
    classes_by_proto_attr_name: Dict[str, Type["SavedDatasetStorage"]] = {}

    def __new__(cls, name, bases, dct):
        kls = type.__new__(cls, name, bases, dct)
        if dct.get("_proto_attr_name"):
            cls.classes_by_proto_attr_name[dct["_proto_attr_name"]] = kls
        return kls

Three classes claim the same key:

Since the assignment is unconditional, the last class whose module is imported wins, and SavedDatasetStorage.from_proto dispatches every custom_storage proto to that one:

@staticmethod
def from_proto(storage_proto: SavedDatasetStorageProto) -> "SavedDatasetStorage":
    proto_attr_name = cast(str, storage_proto.WhichOneof("kind"))
    return _StorageRegistry.classes_by_proto_attr_name[proto_attr_name].from_proto(
        storage_proto
    )

So a saved dataset written against Postgres can be read back as a Clickhouse or Couchbase storage, depending only on which offline store modules happened to be imported first.

The other eight _proto_attr_name values are each claimed exactly once, so custom_storage is the only colliding key.

Scope of what I verified

Read from source: the three claims, the dict assignment, and the from_proto dispatch.

I could not demonstrate the overwrite end to end, because only postgres of the three has its optional dependency installable in my environment — importing the other two fails on clickhouse_connect / couchbase. With postgres alone:

imported SavedDatasetPostgreSQLStorage -> custom_storage now = SavedDatasetPostgreSQLStorage

In practice the collision is latent rather than active, since a deployment rarely installs more than one of these three. That is presumably why it has not surfaced.

Why I am raising it

DataSource.CustomSourceOptions is the only extensible slot in the SavedDatasetStorage oneof:

message SavedDatasetStorage {
  oneof kind {
    DataSource.FileOptions file_storage = 4;
    DataSource.BigQueryOptions bigquery_storage = 5;
    DataSource.RedshiftOptions redshift_storage = 6;
    DataSource.SnowflakeOptions snowflake_storage = 7;
    DataSource.TrinoOptions trino_storage = 8;
    DataSource.SparkOptions spark_storage = 9;
    DataSource.CustomSourceOptions custom_storage = 10;
    DataSource.AthenaOptions athena_storage = 11;
  }
}

So any new CUSTOM_SOURCE-based data source that wants saved-dataset support has to either add a fourth claim to a key that already does not round-trip reliably, or add a new proto field. I ran into this adding a SavedDatasetStorage for a Lance source (#6945) and chose to defer rather than make the collision worse.

Note that DataSource itself does not have this problem: it carries data_source_class_type alongside custom_options and resolves the concrete class from that string, which is exactly the discriminator SavedDatasetStorage lacks.

Possible directions

1. Give SavedDatasetStorage the discriminator DataSource already has, resolving the class from a stored class path instead of from the oneof field name, and keep the registry for backwards compatibility.
2. Give each custom storage its own proto field, which does not scale but is explicit.
3. Make the registry reject a duplicate claim, which would turn this into an import-time error and surface it immediately. That alone is breaking for the three existing classes, so it would need (1) or (2) first.

Happy to send a PR for whichever direction maintainers prefer. I would lean to (1) for symmetry with DataSource.