From 660f7df238e5986b812b02a70c0297fef43a835b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20So=C5=9Bnicki?= <michal.sosnicki@neptune.ai>
Date: Fri, 29 Nov 2024 13:28:25 +0100
Subject: [PATCH] add generate_proto.sh script

---
 .pre-commit-config.yaml                   |  2 +-
 README.md                                 |  9 ++++-
 scripts/{update.sh => generate_client.sh} |  0
 scripts/generate_proto.sh                 | 46 +++++++++++++++++++++++
 scripts/update_descriptor_pool.sh         | 14 -------
 5 files changed, 55 insertions(+), 16 deletions(-)
 rename scripts/{update.sh => generate_client.sh} (100%)
 create mode 100755 scripts/generate_proto.sh
 delete mode 100755 scripts/update_descriptor_pool.sh

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index e241cb6..985b54d 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -33,6 +33,6 @@ repos:
           - grpcio-tools
           - types-Pygments
           - types-protobuf
-exclude: "^src/neptune_api/proto/"
+exclude: "^src/(neptune_api|neptune_retrieval_api)/proto/"
 default_language_version:
     python: python3
diff --git a/README.md b/README.md
index c6b7dc7..820b7b6 100644
--- a/README.md
+++ b/README.md
@@ -78,5 +78,12 @@ Things to know:
 
 Run the following command to regenerate the OpenAPI client:
 ```bash
-scripts/update.sh https://<deployment>.neptune.ai
+scripts/generate_client.sh https://<deployment>.neptune.ai
+```
+
+## Update protobuf model
+
+Generating the protobuf model requires access to the proto files. They are not available publicly as of now.
+```bash
+scripts/generate_proto.sh "$HOME/Projects/neptune"
 ```
diff --git a/scripts/update.sh b/scripts/generate_client.sh
similarity index 100%
rename from scripts/update.sh
rename to scripts/generate_client.sh
diff --git a/scripts/generate_proto.sh b/scripts/generate_proto.sh
new file mode 100755
index 0000000..50908c7
--- /dev/null
+++ b/scripts/generate_proto.sh
@@ -0,0 +1,46 @@
+#!/usr/bin/env bash
+
+set -euxo pipefail
+
+SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )"
+SRC_DIR="${SCRIPT_DIR}/../src"
+NEPTUNE_PROJECT_DIR="${1:-}"
+
+if [ -z "${NEPTUNE_PROJECT_DIR}" ]; then
+    echo "Usage: $0 <neptune_project_dir>"
+    exit 1
+fi
+
+python -m grpc_tools.protoc \
+  --python_out="${SRC_DIR}/neptune_api/proto" \
+  --mypy_out="${SRC_DIR}/neptune_api/proto" \
+  --proto_path="${NEPTUNE_PROJECT_DIR}/libs/models/leaderboard-proto/src/main/proto" \
+  "${NEPTUNE_PROJECT_DIR}/libs/models/leaderboard-proto/src/main/proto/neptune_pb/ingest/v1/"*".proto" \
+  "${NEPTUNE_PROJECT_DIR}/libs/models/leaderboard-proto/src/main/proto/neptune_pb/ingest/v1/pub/"*".proto"
+
+python -m grpc_tools.protoc \
+  --python_out="${SRC_DIR}/neptune_retrieval_api/proto" \
+  --mypy_out="${SRC_DIR}/neptune_retrieval_api/proto" \
+  --proto_path="${NEPTUNE_PROJECT_DIR}/libs/models/leaderboard-proto/src/main/proto" \
+  "${NEPTUNE_PROJECT_DIR}/libs/models/leaderboard-proto/src/main/proto/neptune_pb/api/model/"*".proto"
+
+protol --create-package --in-place \
+  --python-out "${SRC_DIR}/neptune_api/proto" \
+  protoc \
+  --proto-path "${NEPTUNE_PROJECT_DIR}/libs/models/leaderboard-proto/src/main/proto" \
+  "${NEPTUNE_PROJECT_DIR}/libs/models/leaderboard-proto/src/main/proto/neptune_pb/ingest/v1/"*".proto" \
+  "${NEPTUNE_PROJECT_DIR}/libs/models/leaderboard-proto/src/main/proto/neptune_pb/ingest/v1/pub/"*".proto"
+
+protol --create-package --in-place \
+  --python-out "${SRC_DIR}/neptune_retrieval_api/proto" \
+  protoc \
+  --proto-path "${NEPTUNE_PROJECT_DIR}/libs/models/leaderboard-proto/src/main/proto" \
+  "${NEPTUNE_PROJECT_DIR}/libs/models/leaderboard-proto/src/main/proto/neptune_pb/api/model/"*".proto"
+
+for file in $(find "${SRC_DIR}/neptune_api/proto" -name "*.py" -not -name 'descriptor_pool.py'); do
+    sed -i "" "s/from google.protobuf import descriptor_pool/from neptune_api.proto import descriptor_pool/" "${file}"
+done
+
+for file in $(find "${SRC_DIR}/neptune_retrieval_api/proto" -name "*.py" -not -name 'descriptor_pool.py'); do
+    sed -i "" "s/from google.protobuf import descriptor_pool/from neptune_retrieval_api.proto import descriptor_pool/" "${file}"
+done
diff --git a/scripts/update_descriptor_pool.sh b/scripts/update_descriptor_pool.sh
deleted file mode 100755
index 41c8b42..0000000
--- a/scripts/update_descriptor_pool.sh
+++ /dev/null
@@ -1,14 +0,0 @@
-#!/usr/bin/env bash
-
-set -euxo pipefail
-
-SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )"
-SRC_DIR="${SCRIPT_DIR}/../src"
-
-for file in $(find "${SRC_DIR}/neptune_api/proto" -name "*.py"); do
-    sed -i "" "s/from google.protobuf import descriptor_pool/from neptune_api.proto import descriptor_pool/" "${file}"
-done
-
-for file in $(find "${SRC_DIR}/neptune_retrieval_api/proto" -name "*.py"); do
-    sed -i "" "s/from google.protobuf import descriptor_pool/from neptune_retrieval_api.proto import descriptor_pool/" "${file}"
-done