From bfbe1ade90dceae4072a1f44047e34a856ecfd45 Mon Sep 17 00:00:00 2001
From: Lei Wang <jiede.wl@alibaba-inc.com>
Date: Wed, 4 Sep 2024 15:44:22 +0800
Subject: [PATCH] support submit pgql as rgmapping config

Signed-off-by: Lei Wang <jiede.wl@alibaba-inc.com>
---
 Dockerfile            |  5 ++++-
 scripts/controller.py | 40 ++++++++++++++++++++++++++++++++++++++++
 scripts/gart_cli.py   | 24 ++++++++++++++++++++++++
 3 files changed, 68 insertions(+), 1 deletion(-)

diff --git a/Dockerfile b/Dockerfile
index 5583a7c..e8f58b2 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -69,10 +69,13 @@ RUN if [ "$build_type" = "Controller" ]; then \
   curl -LO "https://dl.k8s.io/release/$(curl -L -s https://dl.k8s.io/release/stable.txt)/bin/linux/amd64/kubectl" && \
   chmod +x ./kubectl && \
   mv ./kubectl /usr/local/bin/kubectl && \
-  apt-get update && apt-get install -y openmpi-bin libopenmpi-dev && \
+  apt-get update && apt-get install -y openmpi-bin libopenmpi-dev maven && \
   rm -rf /var/lib/apt/lists/* && \
   pip3 install tenacity==8.3.0 && \
   pip3 install flask kubernetes etcd3; \
+  git clone https://github.com/oracle/pgql-lang.git; \
+  cd pgql-lang; \
+  sh install.sh; \
   fi
 
 # Find the Kafka directory and write its path to a file
diff --git a/scripts/controller.py b/scripts/controller.py
index 456f7ce..fdc4094 100755
--- a/scripts/controller.py
+++ b/scripts/controller.py
@@ -46,6 +46,46 @@ def submit_config():
         return jsonify({"error": str(e)}), 400
 
 
+@app.route("/submit-pgql-config", methods=["POST"])
+def submit_pgql_config():
+    if "file" not in request.files:
+        return jsonify({"error": "No file part in the request"}), 400
+    file = request.files["file"]
+    if file.filename == "":
+        return jsonify({"error": "No selected file"}), 400
+
+    try:
+        content = file.read()
+        with open("/tmp/pgql_config.sql", "wb") as f:
+            f.write(content)
+        subprocess.run(
+            [
+                "/bin/bash",
+                "-c",
+                "cd /workspace/gart/pgql/ && ./run.sh sql2yaml /tmp/pgql_config.sql /tmp/pgql_config.yaml",
+            ]
+        )
+        with open("/tmp/pgql_config.yaml", "r") as f:
+            yaml_content = f.read()
+        yaml_content = yaml_content.replace("!!gart.pgql.GSchema", "")
+        etcd_server = os.getenv("ETCD_SERVICE", "etcd")
+        if not etcd_server.startswith(("http://", "https://")):
+            etcd_server = f"http://{etcd_server}"
+        etcd_prefix = os.getenv("ETCD_PREFIX", "gart_meta_")
+        etcd_host = etcd_server.split("://")[1].split(":")[0]
+        etcd_port = etcd_server.split(":")[2]
+        etcd_client = etcd3.client(host=etcd_host, port=etcd_port)
+        while True:
+            try:
+                etcd_client.put(etcd_prefix + "gart_rg_mapping_yaml", yaml_content)
+                break
+            except Exception as e:
+                time.sleep(5)
+        return "PGQL config submitted", 200
+    except Exception as e:
+        return jsonify({"error": str(e)}), 400
+
+
 @app.route("/control/pause", methods=["POST"])
 def pause():
     subprocess.run(
diff --git a/scripts/gart_cli.py b/scripts/gart_cli.py
index 2823580..877781c 100755
--- a/scripts/gart_cli.py
+++ b/scripts/gart_cli.py
@@ -137,6 +137,30 @@ def submit_config(ctx, config_path):
             click.echo(f"Failed to submit the configuration file: {e}")
 
 
+@cli.command()
+@click.pass_context
+@click.argument("config_path", type=click.Path(exists=True))
+def submit_pgql_config(ctx, config_path):
+    """Submit a new pgql configuration file."""
+    endpoint = ctx.obj.get("endpoint")
+    if not endpoint:
+        click.echo('Please connect to an endpoint first using the "connect" command.')
+        return
+
+    with open(config_path, "rb") as file:
+        files = {"file": (config_path, file)}
+        try:
+            response = requests.post(f"{endpoint}/submit-pgql-config", files=files)
+            response.raise_for_status()
+            click.echo(f"Success: Server responded with {response.status_code} status")
+        except requests.exceptions.HTTPError as e:
+            click.echo(f"Failed to submit the pgql configuration file: {e}")
+        except requests.exceptions.RequestException as e:
+            click.echo(f"Failed to submit the pgql configuration file: {e}")
+        except Exception as e:
+            click.echo(f"Failed to submit the pgql configuration file: {e}")
+
+
 @cli.command()
 @click.pass_context
 @click.argument("algorithm_name")