Protobuf definition fixes, update generate script

mitdbg · Sep 26, 2024 · b8200fb · b8200fb
1 parent 1a8ca92
commit b8200fb
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 9 deletions.
diff --git a/proto/interface/blueprint.proto b/proto/interface/blueprint.proto
@@ -24,11 +24,16 @@ message PhysicalSnapshot {
 
 // Used to indicate the location of data.
 enum Engine {
-  UNKNOWN = 0;
+  ENGINE_UNKNOWN = 0;
   // The data is stored in Aurora.
-  AURORA = 1;
+  ENGINE_AURORA = 1;
   // The data is stored in Redshift.
-  REDSHIFT = 2;
+  ENGINE_REDSHIFT = 2;
   // The data is stored on S3 in an Iceberg table and registered with Athena.
-  ATHENA = 3;
+  ENGINE_ATHENA = 3;
+}
+
+message Provisioning {
+  string instance_type = 1;
+  uint32 num_nodes = 2;
 }
diff --git a/proto/interface/vdbe.proto b/proto/interface/vdbe.proto
@@ -1,11 +1,14 @@
 syntax = "proto3";
 
+package vdbe;
+
 message VirtualEngine {
   string name = 1;
   QueryInterface qi = 2;
   repeated VirtualTable tables = 3;
   uint64 max_staleness_ms = 4;
   // TODO: Meaningful representation for performance properties.
+  // TODO: Flag set for feature support.
 }
 
 message VirtualTable {
@@ -15,13 +18,13 @@ message VirtualTable {
 }
 
 enum QueryInterface {
-  UNKNOWN = 0;
+  QI_UNKNOWN = 0;
 
   // 1xx - Open source SQL DBMSes.
-  SQL_POSTGRESQL = 101;
-  SQL_MYSQL = 102;
+  QI_SQL_POSTGRESQL = 101;
+  QI_SQL_MYSQL = 102;
 
   // 2xx - AWS SQL DBMSes.
-  SQL_AWS_REDSHIFT = 201;
-  SQL_AWS_ATHENA = 202;
+  QI_SQL_AWS_REDSHIFT = 201;
+  QI_SQL_AWS_ATHENA = 202;
 }
diff --git a/tools/generate_proto.sh b/tools/generate_proto.sh
@@ -4,6 +4,20 @@ SCRIPT_PATH=$(cd $(dirname $0) && pwd -P)
 cd $SCRIPT_PATH
 source shared.sh
 
+# NOTE: You need to manually add `# type: ignore` to the .pyi stubs, since the
+# generated code causes an error in mypy's type check.
+python3 -m grpc_tools.protoc \
+  -I../proto \
+  --python_out=../src/brad/proto_gen \
+  --pyi_out=../src/brad/proto_gen \
+  ../proto/interface.proto \
+  ../proto/interface/blueprint.proto \
+  ../proto/interface/schema.proto \
+  ../proto/interface/vdbe.proto
+
+# Fix the import path.
+sed -i -e "s/from interface import/from brad.proto_gen.interface import/g" ../src/brad/proto_gen/interface_pb2.py
+
 python3 -m grpc_tools.protoc \
   -I../proto \
   --python_out=../src/brad/proto_gen \