def _input_artifacts(self, pipeline_name: Text,
input_artifacts: List[Artifact]) -> Channel:
"""Publish input artifacts for test to MLMD and return channel to them."""
connection_config = metadata_store_pb2.ConnectionConfig(
mysql=metadata_store_pb2.MySQLDatabaseConfig(
host='127.0.0.1',
port=3306,
database=self._get_mlmd_db_name(pipeline_name),
user='******',
password=''))
dummy_artifact = (input_artifacts[0].type_name, self._random_id())
output_key = 'dummy_output_%s_%s' % dummy_artifact
producer_component_id = 'dummy_producer_id_%s_%s' % dummy_artifact
producer_component_type = 'dummy_producer_type_%s_%s' % dummy_artifact
# Input artifacts must have a unique name and producer in MLMD.
for artifact in input_artifacts:
artifact.name = output_key
artifact.pipeline_name = pipeline_name
artifact.producer_component = producer_component_id
with metadata.Metadata(connection_config=connection_config) as m:
# Register a dummy execution to metadata store as producer execution.
execution_id = m.register_execution(
exec_properties={},
pipeline_info=data_types.PipelineInfo(
pipeline_name=pipeline_name,
pipeline_root='/dummy_pipeline_root',
# test_utils uses pipeline_name as fixed WORKFLOW_ID.
run_id=pipeline_name,
),
component_info=data_types.ComponentInfo(
component_type=producer_component_type,
component_id=producer_component_id))
# Publish the test input artifact from the dummy execution.
published_artifacts = m.publish_execution(
execution_id=execution_id,
input_dict={},
output_dict={output_key: input_artifacts})
return channel_utils.as_channel(published_artifacts[output_key])
def mysql_metadata_connection_config(
host: Text, port: int, database: Text, username: Text,
password: Text) -> metadata_store_pb2.ConnectionConfig:
"""Convenience function to create mysql-based metadata connection config.
Args:
host: The name or network address of the instance of MySQL to connect to.
port: The port MySQL is using to listen for connections.
database: The name of the database to use.
username: The MySQL login account being used.
password: The password for the MySQL account being used.
Returns:
A metadata_store_pb2.ConnectionConfig based on given metadata db uri.
"""
return metadata_store_pb2.ConnectionConfig(
mysql=metadata_store_pb2.MySQLDatabaseConfig(host=host,
port=port,
database=database,
user=username,
password=password))
