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))