def __init__(self,
config: Union[metadata_store_pb2.ConnectionConfig,
metadata_store_pb2.MetadataStoreClientConfig],
enable_upgrade_migration: bool = False):
"""Initialize the MetadataStore.
MetadataStore can directly connect to either the metadata database or
the metadata store server.
Args:
config: metadata_store_pb2.ConnectionConfig or
metadata_store_pb2.MetadataStoreClientConfig. Configuration to
connect to the database or the metadata store server.
enable_upgrade_migration: if set to True, the library upgrades the db
schema and migrates all data if it connects to an old version backend.
It is ignored when using GRPC client connection config.
"""
self._max_num_retries = 5
if isinstance(config, metadata_store_pb2.ConnectionConfig):
self._using_db_connection = True
migration_options = metadata_store_pb2.MigrationOptions()
migration_options.enable_upgrade_migration = enable_upgrade_migration
self._metadata_store = metadata_store_serialized.CreateMetadataStore(
config.SerializeToString(),
migration_options.SerializeToString())
logging.log(logging.INFO,
'MetadataStore with DB connection initialized')
logging.log(logging.DEBUG, 'ConnectionConfig: %s', config)
if config.HasField('retry_options'):
self._max_num_retries = config.retry_options.max_num_retries
logging.log(
logging.INFO,
'retry options is overwritten: max_num_retries = %d',
self._max_num_retries)
return
if not isinstance(config,
metadata_store_pb2.MetadataStoreClientConfig):
raise ValueError('MetadataStore is expecting either '
'metadata_store_pb2.ConnectionConfig or '
'metadata_store_pb2.MetadataStoreClientConfig')
self._grpc_timeout_sec = None
self._using_db_connection = False
if enable_upgrade_migration:
raise ValueError(
'Upgrade migration is not allowed when using gRPC '
'connection client. Upgrade needs to be performed on '
'the server side.')
channel = self._get_channel(config)
self._metadata_store_stub = (
metadata_store_service_pb2_grpc.MetadataStoreServiceStub(channel))
logging.log(logging.INFO,
'MetadataStore with gRPC connection initialized')
logging.log(logging.DEBUG, 'ConnectionConfig: %s', config)
def __init__(self,
config: Union[metadata_store_pb2.ConnectionConfig,
metadata_store_pb2.MetadataStoreClientConfig],
enable_upgrade_migration: bool = False):
"""Initialize the MetadataStore.
MetadataStore can directly connect to either the metadata database or
the metadata store server.
Args:
config: metadata_store_pb2.ConnectionConfig or
metadata_store_pb2.MetadataStoreClientConfig. Configuration to
connect to the database or the metadata store server.
enable_upgrade_migration: if set to True, the library upgrades the db
schema and migrates all data if it connects to an old version backend.
It is ignored when using GRPC client connection config.
"""
if isinstance(config, metadata_store_pb2.ConnectionConfig):
self._using_db_connection = True
migration_options = metadata_store_pb2.MigrationOptions()
migration_options.enable_upgrade_migration = enable_upgrade_migration
self._metadata_store = metadata_store_serialized.CreateMetadataStore(
config.SerializeToString(),
migration_options.SerializeToString())
logging.log(logging.INFO,
'MetadataStore with DB connection initialized')
return
if not isinstance(config,
metadata_store_pb2.MetadataStoreClientConfig):
raise ValueError('MetadataStore is expecting either'
'metadata_store_pb2.ConnectionConfig or'
'metadata_store_pb2.MetadataStoreClientConfig')
self._using_db_connection = False
if enable_upgrade_migration:
raise ValueError(
'Upgrade migration is not allowed when using gRPC '
'connection client. Upgrade needs to be performed on '
'the server side.')
target = ':'.join([config.host, str(config.port)])
channel = self._get_channel(config, target)
self._metadata_store_stub = (
metadata_store_service_pb2_grpc.MetadataStoreServiceStub(channel))
logging.log(logging.INFO,
'MetadataStore with gRPC connection initialized')
def __init__(self,
config: Union[metadata_store_pb2.ConnectionConfig,
metadata_store_pb2.MetadataStoreClientConfig],
disable_upgrade_migration: bool = False):
"""Initialize the MetadataStore.
MetadataStore can directly connect to either the metadata database or
the metadata store server.
Args:
config: metadata_store_pb2.ConnectionConfig or
metadata_store_pb2.MetadataStoreClientConfig. Configuration to
connect to the database or the metadata store server.
disable_upgrade_migration: if set to True, the library does not upgrades
the db schema and migrates all data if it connects to an old version
backend.
"""
if isinstance(config, metadata_store_pb2.ConnectionConfig):
self._using_db_connection = True
migration_options = metadata_store_pb2.MigrationOptions()
migration_options.disable_upgrade_migration = disable_upgrade_migration
self._metadata_store = metadata_store_serialized.CreateMetadataStore(
config.SerializeToString(),
migration_options.SerializeToString())
# If you remove this line, errors are not thrown correctly.
logging.log(logging.INFO,
'MetadataStore with DB connection initialized')
return
if not isinstance(config,
metadata_store_pb2.MetadataStoreClientConfig):
raise ValueError('MetadataStore is expecting either'
'metadata_store_pb2.ConnectionConfig or'
'metadata_store_pb2.MetadataStoreClientConfig')
self._using_db_connection = False
target = ':'.join([config.host, str(config.port)])
channel = self._get_channel(config, target)
self._metadata_store_stub = (
metadata_store_service_pb2_grpc.MetadataStoreServiceStub(channel))
logging.log(logging.INFO,
'MetadataStore with gRPC connection initialized')
def downgrade_schema(config: metadata_store_pb2.ConnectionConfig,
downgrade_to_schema_version: int) -> None:
"""Downgrades the db specified in the connection config to a schema version.
If `downgrade_to_schema_version` is greater or equals to zero and less than
the current library's schema version, it runs a downgrade transaction to
revert the db schema and migrate the data. The update is transactional, and
any failure will cause a full rollback of the downgrade. Once the downgrade
is done, the user needs to use the older version of the library to connect to
the database.
Args:
config: a metadata_store_pb2.ConnectionConfig having the connection params
downgrade_to_schema_version: downgrades the given database to a specific
version. For v0.13.2 release, the schema_version is 0. For 0.14.0 and
0.15.0 release, the schema_version is 4. More details are described in
g3doc/get_start.md#upgrade-mlmd-library
Raises:
InvalidArgumentError: if the `downgrade_to_schema_version` is not given or
it is negative or greater than the library version.
RuntimeError: if the downgrade is not finished, return detailed error.
"""
if downgrade_to_schema_version < 0:
raise _make_exception('downgrade_to_schema_version not specified',
errors.INVALID_ARGUMENT)
try:
migration_options = metadata_store_pb2.MigrationOptions()
migration_options.disable_upgrade_migration = True
migration_options.downgrade_to_schema_version = downgrade_to_schema_version
metadata_store_serialized.CreateMetadataStore(
config.SerializeToString(), migration_options.SerializeToString())
except RuntimeError as e:
if str(e).startswith('MLMD cannot be downgraded to schema_version'):
raise _make_exception(str(e), errors.INVALID_ARGUMENT)
if not str(e).startswith('Downgrade migration was performed.'):
raise e
# downgrade is done.
logging.log(logging.INFO, str(e))
def __init__(self, config: metadata_store_pb2.ConnectionConfig):
self._metadata_store = metadata_store_serialized.CreateMetadataStore(
config.SerializeToString())
# If you remove this line, errors are not thrown correctly.
logging.log(logging.INFO, "MetadataStore initialized")
