diff --git a/ml_metadata/metadata_store/metadata_store.py b/ml_metadata/metadata_store/metadata_store.py
index 359d30b9e..401246d20 100644
--- a/ml_metadata/metadata_store/metadata_store.py
+++ b/ml_metadata/metadata_store/metadata_store.py
@@ -107,6 +107,7 @@ def __init__(self, config, enable_upgrade_migration: bool = False):
         schema and migrates all data if it connects to an old version backend.
         It is ignored when using gRPC `proto.MetadataStoreClientConfig`.
     """
+    self._config = config
     self._max_num_retries = 5
     self._service_client_wrapper = None
     if isinstance(config, proto.ConnectionConfig):
@@ -139,6 +140,7 @@ def __init__(self, config, enable_upgrade_migration: bool = False):
     logging.log(logging.INFO, 'MetadataStore with gRPC connection initialized')
     logging.log(logging.DEBUG, 'ConnectionConfig: %s', config)
 
+
   def _get_channel(self, config: proto.MetadataStoreClientConfig):
     """Configures the channel, which could be secure or insecure.