feat: provide sqlite configs as cmdline flags

Signed-off-by: Andrea Lamparelli <a.lamparelli95@gmail.com>

Author: lampajr

ml_metadata/metadata_store/metadata_store_server_main.cc

@@ -255,6 +255,52 @@ bool ParsePostgreSQLFlagsBasedServerConfigOrDie(
   return true;
 }
 
+// Converts flag sqlite_connection_mode value to actual enum
+// ConnectionMode. Returns error if conversion is not successful.
+absl::StatusOr<ml_metadata::SqliteMetadataSourceConfig::ConnectionMode> ConvertToConnectionMode(
+    std::string sqlite_connection_mode) {
+  if (sqlite_connection_mode == "READONLY") {
+    return ml_metadata::SqliteMetadataSourceConfig::READONLY;
+  } else if (sqlite_connection_mode == "READWRITE") {
+    return ml_metadata::SqliteMetadataSourceConfig::READWRITE;
+  } else if (sqlite_connection_mode == "READWRITE_OPENCREATE") {
+    return ml_metadata::SqliteMetadataSourceConfig::READWRITE_OPENCREATE;
+  }
+
+  return absl::InvalidArgumentError(
+      "sqlite_connection_mode is not valid, provide value in one of the "
+      "followings: READONLY, READWRITE, READWRITE_OPENCREATE.");
+}
+
+// Returns true if passed parameters were used to construct sqlite connection
+// config and set it to service_config. If connection config
+// construction is not successful, return false.
+bool ParseSQLiteFlagsBasedServerConfigOrDie(
+    const std::string& filename_uri,
+    const std::string& connection_mode,
+    ml_metadata::MetadataStoreServerConfig* server_config) {
+  CHECK(!filename_uri.empty())
+      << "To use sqlite store, all of --filename_uri, "
+         "needs to be provided";
+
+  ml_metadata::ConnectionConfig* connection_config =
+      server_config->mutable_connection_config();
+  ml_metadata::SqliteMetadataSourceConfig* config = connection_config->mutable_sqlite();
+  config->set_filename_uri(filename_uri);
+
+  if (!connection_mode.empty()) {
+    absl::StatusOr<ml_metadata::SqliteMetadataSourceConfig::ConnectionMode> sqlite_connection_mode =
+    ConvertToConnectionMode(connection_mode);
+    CHECK(sqlite_connection_mode.ok())
+      << "sqlite_connection_mode is invalid: "
+      << sqlite_connection_mode.status().message();
+
+    config->set_connection_mode(sqlite_connection_mode.value());
+  }
+
+  return true;
+}
+
 }  // namespace
 
 // gRPC server options
@@ -270,6 +316,7 @@ enum class SourceConfigType {
   kConfigFile,
   kMySql,
   kPostgreSql,
+  kSqlite
 };
 
 // Converts flag metadata_source_config_type value to actual enum
@@ -284,6 +331,8 @@ absl::StatusOr<SourceConfigType> ConvertToSourceConfig(
     return SourceConfigType::kMySql;
   } else if (metadata_source_config_type == "postgresql") {
     return SourceConfigType::kPostgreSql;
+  } else if (metadata_source_config_type == "sqlite") {
+    return SourceConfigType::kSqlite;
   }
 
   return absl::InvalidArgumentError(
@@ -383,6 +432,14 @@ DEFINE_int64(downgrade_db_schema_version, -1,
              "schema version is downgraded to the set value during "
              "initialization(Optional Parameter)");
 
+// SQLite config command line options
+DEFINE_string(sqlite_config_filename_uri, "",
+              "SQLite database filename uri to be used.");
+DEFINE_string(sqlite_config_connection_mode, "",
+              "SQLite connection mode. Possible values are READONLY,"
+              "READWRITE and READWRITE_OPENCREATE, if not specifified "
+              "the default is READWRITE_OPENCREATE.");
+
 // Default connection option for metadata source. It will check for
 // the existence of config file first, and check for mysql flags if
 // config file doesn't exist. Otherwise, it will create fake database.
@@ -494,6 +551,35 @@ BuildPostgreSQLConnectionConfig() {
   }
 }
 
+// Constructs Connection Config for SQLite database. Requires to
+// set metadata_source_config_type as "sqlite", then provide necessary
+// information in flags that have prefix of `sqlite_`.
+// Example run:
+// sudo docker run --name "${MLMD_GRPC_CONTAINER}" \
+//   -p ${MLMD_GRPC_PORT}:${MLMD_GRPC_PORT} \
+//   --network="${GRPC_E2E_BRIDGE_NETWORK}"\
+//   --entrypoint /bin/metadata_store_server -d "${MLMD_DOCKER_IMAGE}" \
+//   --grpc_port=${MLMD_GRPC_PORT} \
+//   --metadata_source_config_type="sqlite" \
+//   --sqlite_config_filename_uri=${MLMD_SQLITE_DB_URI} \
+//   --sqlite_config_connection_mode=READWRITE_OPENCREATE
+// @return server configuration that contains connection config set by
+// sqlite flags. Or error status if failed.
+absl::StatusOr<ml_metadata::MetadataStoreServerConfig>
+BuildSQLiteConnectionConfig() {
+  ml_metadata::MetadataStoreServerConfig server_config;
+  if (ParseSQLiteFlagsBasedServerConfigOrDie(
+          (FLAGS_sqlite_config_filename_uri),
+          (FLAGS_sqlite_config_connection_mode),
+          &server_config)) {
+    return server_config;
+  } else {
+    LOG(ERROR) << "Unable to construct server config using sqlite flags.";
+    return absl::InvalidArgumentError(
+        "Unable to construct server config using sqlite flags.");
+  }
+}
+
 int main(int argc, char** argv) {
   gflags::ParseCommandLineFlags(&argc, &argv, true);
 
@@ -521,6 +607,8 @@ int main(int argc, char** argv) {
     server_config_status = BuildMySQLConnectionConfig();
   } else if (source_config.value() == SourceConfigType::kPostgreSql) {
     server_config_status = BuildPostgreSQLConnectionConfig();
+  } else if (source_config.value() == SourceConfigType::kSqlite) {
+    server_config_status = BuildSQLiteConnectionConfig();
   } else {
     LOG(ERROR) << "metadata_source_config_type is invalid: "
                << metadata_source_config_type;