From a8b5fe77fa0b22a4e94d52933059bfae92f05e97 Mon Sep 17 00:00:00 2001
From: Xiang Fu <xiangfu.1024@gmail.com>
Date: Wed, 13 Sep 2023 20:08:02 -0700
Subject: [PATCH] Allow override/force options when add schema (#11572)

---
 .../tools/admin/command/AddSchemaCommand.java | 29 ++++++++++++++++---
 1 file changed, 25 insertions(+), 4 deletions(-)

diff --git a/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/AddSchemaCommand.java b/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/AddSchemaCommand.java
index 18d1dfd6c457..7ddc21b88dee 100644
--- a/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/AddSchemaCommand.java
+++ b/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/AddSchemaCommand.java
@@ -20,6 +20,7 @@
 
 import java.io.File;
 import java.io.FileNotFoundException;
+import java.net.URI;
 import java.util.Collections;
 import org.apache.pinot.common.utils.FileUploadDownloadClient;
 import org.apache.pinot.spi.auth.AuthProvider;
@@ -48,6 +49,14 @@ public class AddSchemaCommand extends AbstractBaseAdminCommand implements Comman
   @CommandLine.Option(names = {"-schemaFile"}, required = true, description = "Path to schema file.")
   private String _schemaFile = null;
 
+  @CommandLine.Option(names = {"-override"}, required = false, description = "Whether to override the schema if the "
+      + "schema exists.")
+  private boolean _override = false;
+
+  @CommandLine.Option(names = {"-force"}, required = false, description = "Whether to force overriding the schema if "
+      + "the schema exists.")
+  private boolean _force = false;
+
   @CommandLine.Option(names = {"-exec"}, required = false, description = "Execute the command.")
   private boolean _exec;
 
@@ -87,8 +96,8 @@ public String getName() {
   @Override
   public String toString() {
     String retString = ("AddSchema -controllerProtocol " + _controllerProtocol + " -controllerHost " + _controllerHost
-        + " -controllerPort " + _controllerPort + " -schemaFile " + _schemaFile + " -user " + _user + " -password "
-        + "[hidden]");
+        + " -controllerPort " + _controllerPort + " -schemaFile " + _schemaFile + " -override " + _override + " _force "
+        + _force + " -user " + _user + " -password " + "[hidden]");
 
     return ((_exec) ? (retString + " -exec") : retString);
   }
@@ -117,6 +126,16 @@ public AddSchemaCommand setSchemaFilePath(String schemaFilePath) {
     return this;
   }
 
+  public AddSchemaCommand setOverride(boolean override) {
+    _override = override;
+    return this;
+  }
+
+  public AddSchemaCommand setForce(boolean force) {
+    _force = force;
+    return this;
+  }
+
   public void setUser(String user) {
     _user = user;
   }
@@ -155,8 +174,10 @@ public boolean execute()
 
     Schema schema = Schema.fromFile(schemaFile);
     try (FileUploadDownloadClient fileUploadDownloadClient = new FileUploadDownloadClient()) {
-      fileUploadDownloadClient.addSchema(FileUploadDownloadClient
-              .getUploadSchemaURI(_controllerProtocol, _controllerHost, Integer.parseInt(_controllerPort)),
+      URI schemaURI = FileUploadDownloadClient
+          .getUploadSchemaURI(_controllerProtocol, _controllerHost, Integer.parseInt(_controllerPort));
+      schemaURI = new URI(schemaURI + "?override=" + _override + "?force=" + _force);
+      fileUploadDownloadClient.addSchema(schemaURI,
           schema.getSchemaName(), schemaFile, makeAuthHeaders(makeAuthProvider(_authProvider, _authTokenUrl, _authToken,
               _user, _password)), Collections.emptyList());
     } catch (Exception e) {