Supports SSL options for MySQL metadata source. Fixes github issue #20.

PiperOrigin-RevId: 279982197

RELEASE.md

@@ -5,6 +5,7 @@
 *   Add migration options to gRPC MetadataStoreServerConfig.
 *   Disable auto schema migration by default during connection. The user needs
     to explicitly enable it when connecting an older database.
+*   Support SSL options when using MySQL metadata source.
 
 ## Bug Fixes and Other Changes
 

ml_metadata/metadata_store/mysql_metadata_source.cc

@@ -107,7 +107,20 @@ Status MySqlMetadataSource::ConnectImpl() {
   }
 
   // Explicitly setup the thread-local initializer.
-  TF_RETURN_IF_ERROR(ThreadInitAccess());
+  TF_RETURN_WITH_CONTEXT_IF_ERROR(ThreadInitAccess(),
+                                  "MySql thread init failed at ConnectImpl");
+
+  // Set connection options
+  if (config_.has_ssl_options()) {
+    const MySQLDatabaseConfig::SSLOptions& ssl = config_.ssl_options();
+    // The method set mysql_options, and always return 0. The connection options
+    // are used in the `mysql_real_connect`.
+    mysql_ssl_set(db_, ssl.key().empty() ? nullptr : ssl.key().c_str(),
+                  ssl.cert().empty() ? nullptr : ssl.cert().c_str(),
+                  ssl.ca().empty() ? nullptr : ssl.ca().c_str(),
+                  ssl.capath().empty() ? nullptr : ssl.capath().c_str(),
+                  ssl.cipher().empty() ? nullptr : ssl.cipher().c_str());
+  }
 
   // Connect to the MYSQL server.
   if (!mysql_real_connect(
@@ -152,7 +165,8 @@ Status MySqlMetadataSource::CloseImpl() {
 
 Status MySqlMetadataSource::ExecuteQueryImpl(const string& query,
                                              RecordSet* results) {
-  TF_RETURN_WITH_CONTEXT_IF_ERROR(ThreadInitAccess(), "ExecuteQueryImpl");
+  TF_RETURN_WITH_CONTEXT_IF_ERROR(
+      ThreadInitAccess(), "MySql thread init failed at ExecuteQueryImpl");
 
   // Run the query.
   Status status = RunQuery(query);
@@ -169,20 +183,22 @@ Status MySqlMetadataSource::ExecuteQueryImpl(const string& query,
 
 Status MySqlMetadataSource::CommitImpl() {
   constexpr char kCommitTransaction[] = "COMMIT";
-
-  TF_RETURN_WITH_CONTEXT_IF_ERROR(ThreadInitAccess(), "CommitImpl");
+  TF_RETURN_WITH_CONTEXT_IF_ERROR(ThreadInitAccess(),
+                                  "MySql thread init failed at CommitImpl");
   return RunQuery(kCommitTransaction);
 }
 
 Status MySqlMetadataSource::RollbackImpl() {
   constexpr char kRollbackTransaction[] = "ROLLBACK";
-
-  TF_RETURN_WITH_CONTEXT_IF_ERROR(ThreadInitAccess(), "RollbackImpl");
+  TF_RETURN_WITH_CONTEXT_IF_ERROR(ThreadInitAccess(),
+                                  "MySql thread init failed at RollbackImpl");
   return RunQuery(kRollbackTransaction);
 }
 
 Status MySqlMetadataSource::BeginImpl() {
   constexpr char kBeginTransaction[] = "START TRANSACTION";
+  TF_RETURN_WITH_CONTEXT_IF_ERROR(ThreadInitAccess(),
+                                  "MySql thread init failed at BeginImpl");
   return RunQuery(kBeginTransaction);
 }
 

ml_metadata/proto/metadata_store.proto

@@ -367,6 +367,24 @@ message MySQLDatabaseConfig {
   // The Unix socket to use to connect to the server. If unspecified, a
   // `host` must be provided.
   optional string socket = 6;
+
+  // The options to establish encrypted connections to MySQL using SSL.
+  message SSLOptions {
+    // The path name of the client private key file.
+    optional string key = 1;
+    // The path name of the client public key certificate file.
+    optional string cert = 2;
+    // The path name of the CA certificate file.
+    optional string ca = 3;
+    // The path name of the directory that contains trusted SSL CA certificates.
+    optional string capath = 4;
+    // The list of permissible ciphers for SSL encryption.
+    optional string cipher = 5;
+  }
+  // If the field is set, the ssl options are set in mysql_options before
+  // establishing a connection. It is ignored if the mysql server does not
+  // enable SSL.
+  optional SSLOptions ssl_options = 7;
 }
 
 // A config contains the parameters when using with SqliteMetadatSource.