Add support for protobuf

.github/workflows/build-test.yml

@@ -48,5 +48,11 @@ jobs:
         pip install .
     - name: Test with pytest
       run: |
-        pip install ./tests/snowflake-telemetry-test-utils
+        pip install "./tests/snowflake-telemetry-test-utils[all]"
         pytest
+    - name: Test with tox
+      run: |
+        version="${{ matrix.python-version }}"
+        version_without_dot="${version//./}"
+        pip install tox
+        env TOX_SKIP_ENV='^(?!py'"$version_without_dot"'-)' tox

.github/workflows/check-codegen.yml

@@ -10,13 +10,15 @@ on:
     paths:
       - "scripts/**"
       - "src/snowflake/telemetry/_internal/opentelemetry/proto/**"
+      - "src/snowflake/telemetry/_internal/proto_impl/**"
       - "src/snowflake/telemetry/serialize/**"
       - ".github/workflows/check-codegen.yml"
   pull_request:
     branches: [ "main" ]
     paths:
       - "scripts/**"
       - "src/snowflake/telemetry/_internal/opentelemetry/proto/**"
+      - "src/snowflake/telemetry/_internal/proto_impl/**"
       - "src/snowflake/telemetry/serialize/**"
       - ".github/workflows/check-codegen.yml"
 
@@ -29,10 +31,28 @@ jobs:
       uses: actions/setup-python@v3
       with:
         python-version: "3.11"
-    - name: Run codegen script
+    - name: Run py codegen script
       run: |
         rm -rf src/snowflake/telemetry/_internal/opentelemetry/proto/
+        rm -rf src/snowflake/telemetry/_internal/proto_impl/py/
         ./scripts/proto_codegen.sh
     - name: Check for changes
       run: |
         git diff --exit-code || { echo "Code generation produced changes! Regenerate the code using ./scripts/proto_codegen.sh"; exit 1; }
+  check-codegen-pb:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v3
+    - name: Set up Python
+      uses: actions/setup-python@v3
+      with:
+        python-version: "3.9"
+    - name: Run pb codegen script
+      run: |
+        rm -rf src/snowflake/telemetry/_internal/proto_impl/v3/
+        rm -rf src/snowflake/telemetry/_internal/proto_impl/v5/
+        ./scripts/proto_codegen_vX.sh 3
+        ./scripts/proto_codegen_vX.sh 5
+    - name: Check for changes
+      run: |
+        git diff --exit-code || { echo "Code generation produced changes! Regenerate the code using ./scripts/proto_codegen_vX.sh <version>"; exit 1; }

.github/workflows/check-vendor.yml

@@ -19,7 +19,7 @@ on:
       - ".github/workflows/check-vendor.yml"
 
 jobs:
-  check-codegen:
+  check-vendor:
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v3

README.md

@@ -33,7 +33,7 @@ cd snowflake-telemetry-python
 python3 -m venv .venv
 source .venv/bin/activate
 pip install --upgrade pip
-pip install . ./tests/snowflake-telemetry-test-utils
+pip install . "./tests/snowflake-telemetry-test-utils[all]"
 ```
 
 ### Code generation

scripts/gen-requirements-v3.txt

@@ -0,0 +1,2 @@
+grpcio-tools==1.48.1
+mypy-protobuf~=3.0.0

scripts/gen-requirements-v5.txt

@@ -0,0 +1,2 @@
+grpcio-tools==1.63.2
+mypy-protobuf~=3.5.0

scripts/plugin.py

@@ -21,13 +21,24 @@
     EnumValueDescriptorProto,
     DescriptorProto,
 )
-from jinja2 import Environment, FileSystemLoader
+import jinja2
 import black
 import isort.api
 
 INLINE_OPTIMIZATION = True
-FILE_PATH_PREFIX = "snowflake.telemetry._internal"
-FILE_NAME_SUFFIX = "_marshaler"
+# The location of the pure python implementation of the protobuf messages
+PY_PATH_PREFIX = "snowflake.telemetry._internal.proto_impl.py"
+PY_PATH_RELATIVE = "proto_impl/py"
+# The location of the protobuf v3.X.X and v4.X.X implementation of the protobuf messages
+PB3_PATH_PREFIX = "snowflake.telemetry._internal.proto_impl.v3"
+# The location of the protobuf v5.X.X implementation of the protobuf messages
+PB5_PATH_PREFIX = "snowflake.telemetry._internal.proto_impl.v5"
+
+# Suffixes for the generated files
+PROXY_FILE_SUFFIX = "_proxy"
+PY_FILE_SUFFIX = "_marshaler"
+PB_FILE_SUFFIX = "_pb2"
+
 OPENTELEMETRY_PROTO_DIR = os.environ["OPENTELEMETRY_PROTO_DIR"]
 
 # Inline utility functions
@@ -284,6 +295,9 @@ class FileTemplate:
     imports: List[str] = field(default_factory=list)
     name: str = ""
     preamble: str = ""
+    path_py: str = ""
+    path_pb3: str = ""
+    path_pb5: str = ""
 
     @staticmethod
     def from_descriptor(descriptor: FileDescriptorProto) -> "FileTemplate":
@@ -304,17 +318,41 @@ def from_descriptor(descriptor: FileDescriptorProto) -> "FileTemplate":
             if descriptor.name.startswith(path):
                 continue
             path = path.replace("/", ".")
-            path = f"{FILE_PATH_PREFIX}.{path}{FILE_NAME_SUFFIX}"
+            path = f"{PY_PATH_PREFIX}.{path}{PY_FILE_SUFFIX}"
             imports.append(path)
 
+        path = descriptor.name.replace(".proto", "").replace("/", ".")
+        path_py = f"{PY_PATH_PREFIX}.{path}{PY_FILE_SUFFIX}"
+        path_pb = f"{PB3_PATH_PREFIX}.{path}{PB_FILE_SUFFIX}"
+        path_pb5 = f"{PB5_PATH_PREFIX}.{path}{PB_FILE_SUFFIX}"
+
         return FileTemplate(
             messages=[MessageTemplate.from_descriptor(message) for message in descriptor.message_type],
             enums=[EnumTemplate.from_descriptor(enum) for enum in descriptor.enum_type],
             imports=imports,
             name=descriptor.name,
             preamble=preamble,
+            path_py=path_py,
+            path_pb3=path_pb,
+            path_pb5=path_pb5,
         )
 
+def format_code_str(code: str) -> str:
+    code = isort.api.sort_code_string(
+        code = code,
+        show_diff=False,
+        profile="black",
+        combine_as_imports=True,
+        lines_after_imports=2,
+        quiet=True,
+        force_grid_wrap=2,
+    )
+    code = black.format_str(
+        src_contents=code,
+        mode=black.Mode(),
+    )
+    return code
+
 def main():
     request = plugin.CodeGeneratorRequest()
     request.ParseFromString(sys.stdin.buffer.read())
@@ -324,30 +362,26 @@ def main():
     # https://github.com/protocolbuffers/protobuf/blob/main/docs/implementing_proto3_presence.md
     response.supported_features = plugin.CodeGeneratorResponse.FEATURE_PROTO3_OPTIONAL
 
-    template_env = Environment(loader=FileSystemLoader(f"{os.path.dirname(os.path.realpath(__file__))}/templates"))
+    template_env = jinja2.Environment(loader=jinja2.FileSystemLoader(f"{os.path.dirname(os.path.realpath(__file__))}/templates"))
     jinja_body_template = template_env.get_template("template.py.jinja2")
+    jinja_proxy_template = template_env.get_template("proxy.py.jinja2")
 
     for proto_file in request.proto_file:
-        file_name = re.sub(r"\.proto$", f"{FILE_NAME_SUFFIX}.py", proto_file.name)
-        file_descriptor_proto = proto_file
-
-        file_template = FileTemplate.from_descriptor(file_descriptor_proto)
+        file_template = FileTemplate.from_descriptor(proto_file)
 
+        # Generate the marshaler file
+        file_name = re.sub(r"\.proto$", f"{PY_FILE_SUFFIX}.py", proto_file.name)
+        file_name = f"{PY_PATH_RELATIVE}/{file_name}"
         code = jinja_body_template.render(file_template=file_template)
-        code = isort.api.sort_code_string(
-            code = code,
-            show_diff=False,
-            profile="black",
-            combine_as_imports=True,
-            lines_after_imports=2,
-            quiet=True,
-            force_grid_wrap=2,
-        )
-        code = black.format_str(
-            src_contents=code,
-            mode=black.Mode(),
-        )
+        code = format_code_str(code)
+        response_file = response.file.add()
+        response_file.name = file_name
+        response_file.content = code
 
+        # Generate the proxy file
+        file_name = re.sub(r"\.proto$", f"{PROXY_FILE_SUFFIX}.py", proto_file.name)
+        code = jinja_proxy_template.render(file_template=file_template)
+        code = format_code_str(code)
         response_file = response.file.add()
         response_file.name = file_name
         response_file.content = code

scripts/proto_codegen.sh

@@ -20,6 +20,7 @@ PROTO_REPO_DIR=${PROTO_REPO_DIR:-"/tmp/opentelemetry-proto"}
 # root of opentelemetry-python repo
 repo_root="$(git rev-parse --show-toplevel)"
 venv_dir="/tmp/proto_codegen_venv"
+code_dir="proto_impl/py"
 
 # run on exit even if crash
 cleanup() {
@@ -55,8 +56,13 @@ fi
 cd $repo_root/src/snowflake/telemetry/_internal
 
 # clean up old generated code
+# remove all the proxy files
 mkdir -p opentelemetry/proto
-find opentelemetry/proto/ -regex ".*_marshaler\.py" -exec rm {} +
+find opentelemetry/proto/ -regex ".*_proxy\.py" -exec rm {} +
+
+# remove all the py implementation marshaler files
+mkdir -p $code_dir/opentelemetry/proto
+find $code_dir/opentelemetry/proto/ -regex ".*_marshaler\.py" -exec rm {} +
 
 # generate proto code for all protos
 all_protos=$(find $PROTO_REPO_DIR/ -iname "*.proto")

scripts/proto_codegen_vX.sh

@@ -0,0 +1,87 @@
+#!/bin/bash
+#
+# Regenerate python code from OTLP protos in
+# https://github.com/open-telemetry/opentelemetry-proto
+#
+# To use, update PROTO_REPO_BRANCH_OR_COMMIT variable below to a commit hash or
+# tag in opentelemtry-proto repo that you want to build off of. Then, just run
+# this script to update the proto files. Commit the changes as well as any
+# fixes needed in the OTLP exporter.
+#
+# Optional envars:
+#   PROTO_REPO_DIR - the path to an existing checkout of the opentelemetry-proto repo
+
+# Pinned commit/branch/tag for the current version used in opentelemetry-proto python package.
+PROTO_REPO_BRANCH_OR_COMMIT="v1.2.0"
+
+set -e
+
+PB_VERSION_MAJOR=${1:-"3"}
+PROTO_REPO_DIR=${PROTO_REPO_DIR:-"/tmp/opentelemetry-proto"}
+# root of opentelemetry-python repo
+repo_root="$(git rev-parse --show-toplevel)"
+venv_dir="/tmp/proto_codegen_venv"
+code_dir="proto_impl/v$PB_VERSION_MAJOR"
+
+# run on exit even if crash
+cleanup() {
+    echo "Deleting $venv_dir"
+    rm -rf $venv_dir
+}
+trap cleanup EXIT
+
+echo "Creating temporary virtualenv at $venv_dir using $(python3 --version)"
+python3 -m venv $venv_dir
+source $venv_dir/bin/activate
+pip install --upgrade pip
+python -m pip install --upgrade setuptools
+python -m pip install \
+    -c $repo_root/scripts/gen-requirements-v$PB_VERSION_MAJOR.txt \
+    grpcio-tools mypy-protobuf
+
+echo 'python -m grpc_tools.protoc --version'
+python -m grpc_tools.protoc --version
+
+# Clone the proto repo if it doesn't exist
+if [ ! -d "$PROTO_REPO_DIR" ]; then
+    git clone https://github.com/open-telemetry/opentelemetry-proto.git $PROTO_REPO_DIR
+fi
+
+# Pull in changes and switch to requested branch
+(
+    cd $PROTO_REPO_DIR
+    git fetch --all
+    git checkout $PROTO_REPO_BRANCH_OR_COMMIT
+    # pull if PROTO_REPO_BRANCH_OR_COMMIT is not a detached head
+    git symbolic-ref -q HEAD && git pull --ff-only || true
+)
+
+cd $repo_root/src/snowflake/telemetry/_internal
+
+# clean up old generated code
+mkdir -p $code_dir/opentelemetry/proto
+find $code_dir -regex ".*_pb2\.py" -exec rm {} +
+find $code_dir -regex ".*_pb2\.pyi" -exec rm {} +
+
+# generate proto code for all protos
+all_protos=$(find $PROTO_REPO_DIR/ -iname "*.proto")
+python -m grpc_tools.protoc \
+    -I $PROTO_REPO_DIR \
+    --python_out=./$code_dir \
+    $all_protos
+
+# since we do not have the generated files in the base directory
+# we need to use sed to update the imports
+
+# If MacOS need '' after -i
+# Detect the OS (macOS or Linux)
+if [[ "$OSTYPE" == "darwin"* ]]; then
+  SED_CMD="sed -i ''"
+else
+  SED_CMD="sed -i"
+fi
+
+find $code_dir -type f \( -name "*.py" -o -name "*.pyi" \) -exec $SED_CMD 's/^import \([^ ]*\)_pb2 as \([^ ]*\)$/import snowflake.telemetry._internal.proto_impl.v'"$PB_VERSION_MAJOR"'.\1_pb2 as \2/' {} +
+find $code_dir -type f \( -name "*.py" -o -name "*.pyi" \) -exec $SED_CMD 's/^from \([^ ]*\) import \([^ ]*\)_pb2 as \([^ ]*\)$/from snowflake.telemetry._internal.proto_impl.v'"$PB_VERSION_MAJOR"'.\1 import \2_pb2 as \3/' {} +
+
+find $code_dir -type f \( -name "*.py''" -o -name "*.pyi''" \) -exec rm {} +

scripts/templates/proxy.py.jinja2

@@ -0,0 +1,14 @@
+# Generated by the protoc compiler with a custom plugin. DO NOT EDIT!
+# sources: {{ file_template.name }}
+#
+# Copyright (c) 2012-2024 Snowflake Inc. All rights reserved.
+#
+
+from snowflake.telemetry._internal.proto_proxy import PROTOBUF_VERSION_MAJOR
+
+if PROTOBUF_VERSION_MAJOR == 3 or PROTOBUF_VERSION_MAJOR == 4:
+    from {{ file_template.path_pb3 }} import *
+elif PROTOBUF_VERSION_MAJOR == 5:
+    from {{ file_template.path_pb5 }} import *
+else:
+    from {{ file_template.path_py }} import *

scripts/vendor_otlp_proto_common.sh

@@ -47,12 +47,14 @@ fi
 
 # Replace all the imports strings in the copied python files
 #   opentelemetry.exporter to snowflake.telemetry._internal.opentelemetry.exporter
-#   opentelemetry.proto.*_pb2 to snowflake.telemetry._internal.opentelemetry.proto.*_marshaler
+#   opentelemetry.proto.*_pb2 to snowflake.telemetry._internal.opentelemetry.proto.*_proxy
 
 find opentelemetry/exporter -type f -name "*.py" -exec $SED_CMD 's/opentelemetry.exporter/snowflake.telemetry._internal.opentelemetry.exporter/g' {} +
-find opentelemetry/exporter -type f -name "*.py" -exec $SED_CMD 's/opentelemetry\.proto\(.*\)_pb2/snowflake.telemetry._internal.opentelemetry.proto\1_marshaler/g' {} +
+find opentelemetry/exporter -type f -name "*.py" -exec $SED_CMD 's/opentelemetry\.proto\(.*\)_pb2/snowflake.telemetry._internal.opentelemetry.proto\1_proxy/g' {} +
 
 
 # Add a notice to the top of every file in compliance with Apache 2.0 to indicate that the file has been modified
 # https://www.apache.org/licenses/LICENSE-2.0
 find opentelemetry/exporter -type f -name "*.py" -exec $SED_CMD '14s|^|#\n# This file has been modified from the original source code at\n#\n#     https://github.com/open-telemetry/opentelemetry-python/tree/'"$REPO_BRANCH_OR_COMMIT"'\n#\n# by Snowflake Inc.\n|' {} +
+
+find opentelemetry/exporter -type f -name "*.py''" -exec rm {} +

src/snowflake/telemetry/_internal/exporter/otlp/proto/logs/__init__.py

@@ -24,7 +24,7 @@
 from snowflake.telemetry._internal.opentelemetry.exporter.otlp.proto.common._log_encoder import (
     encode_logs,
 )
-from snowflake.telemetry._internal.opentelemetry.proto.logs.v1.logs_marshaler import LogsData
+from snowflake.telemetry._internal.opentelemetry.proto.logs.v1.logs_proxy import LogsData
 from opentelemetry.sdk.resources import Resource
 from opentelemetry.sdk._logs import export
 from opentelemetry.sdk import _logs

src/snowflake/telemetry/_internal/exporter/otlp/proto/metrics/__init__.py

@@ -20,7 +20,7 @@
 from snowflake.telemetry._internal.opentelemetry.exporter.otlp.proto.common.metrics_encoder import (
     encode_metrics,
 )
-from snowflake.telemetry._internal.opentelemetry.proto.metrics.v1.metrics_marshaler import MetricsData as PB2MetricsData
+from snowflake.telemetry._internal.opentelemetry.proto.metrics.v1.metrics_proxy import MetricsData as PB2MetricsData
 from opentelemetry.sdk.metrics.export import (
     AggregationTemporality,
     MetricExportResult,

src/snowflake/telemetry/_internal/exporter/otlp/proto/traces/__init__.py

@@ -19,7 +19,7 @@
 from snowflake.telemetry._internal.opentelemetry.exporter.otlp.proto.common.trace_encoder import (
     encode_spans,
 )
-from snowflake.telemetry._internal.opentelemetry.proto.trace.v1.trace_marshaler import TracesData
+from snowflake.telemetry._internal.opentelemetry.proto.trace.v1.trace_proxy import TracesData
 from opentelemetry.sdk.trace import ReadableSpan
 from opentelemetry.sdk.trace.export import (
     SpanExportResult,