-
diff --git a/api/buf.gen.aperture-langs.yaml b/api/buf.gen.aperture-langs.yaml
index 196fa990d0..a2bab5d8ec 100644
--- a/api/buf.gen.aperture-langs.yaml
+++ b/api/buf.gen.aperture-langs.yaml
@@ -16,7 +16,7 @@ plugins:
out: gen/proto/java
- plugin: buf.build/grpc/python
out: gen/proto/python
- - plugin: buf.build/protocolbuffers/python:v24.4
+ - plugin: buf.build/protocolbuffers/python:v25.2
out: gen/proto/python
- plugin: buf.build/protocolbuffers/pyi
out: gen/proto/python
diff --git a/concurrency-limiting-test.yaml b/concurrency-limiting-test.yaml
deleted file mode 100644
index dfd4641176..0000000000
--- a/concurrency-limiting-test.yaml
+++ /dev/null
@@ -1,18 +0,0 @@
-blueprint: concurrency-limiting/base
-uri: github.com/fluxninja/aperture/blueprints@latest
-policy:
- components: []
- policy_name: concurrency-limit-test
- resources:
- flow_control:
- classifiers: []
- concurrency_limiter:
- alerter:
- alert_name: "Too many inflight requests"
- max_concurrency: 10
- parameters:
- limit_by_label_key: "user_id"
- max_inflight_duration: 60s
- request_parameters: {}
- selectors:
- - control_point: concurrency-limiting-feature
diff --git a/sdks/aperture-go/examples/manual/go.mod b/sdks/aperture-go/examples/manual/go.mod
index 4cbb9365d3..08ff269574 100644
--- a/sdks/aperture-go/examples/manual/go.mod
+++ b/sdks/aperture-go/examples/manual/go.mod
@@ -12,7 +12,7 @@ require (
require (
github.com/cenkalti/backoff/v4 v4.2.1 // indirect
github.com/envoyproxy/protoc-gen-validate v1.0.2 // indirect
- github.com/fluxninja/aperture/api/v2 v2.0.0-20240123210709-ea3fb20237aa // indirect
+ github.com/fluxninja/aperture/api/v2 v2.0.0-20240126005715-ea93b670518e // indirect
github.com/go-logr/logr v1.3.0 // indirect
github.com/go-logr/stdr v1.2.2 // indirect
github.com/golang/protobuf v1.5.3 // indirect
diff --git a/sdks/aperture-go/examples/manual/go.sum b/sdks/aperture-go/examples/manual/go.sum
index b50f8e2256..e275297f76 100644
--- a/sdks/aperture-go/examples/manual/go.sum
+++ b/sdks/aperture-go/examples/manual/go.sum
@@ -4,8 +4,8 @@ github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/envoyproxy/protoc-gen-validate v1.0.2 h1:QkIBuU5k+x7/QXPvPPnWXWlCdaBFApVqftFV6k087DA=
github.com/envoyproxy/protoc-gen-validate v1.0.2/go.mod h1:GpiZQP3dDbg4JouG/NNS7QWXpgx6x8QiMKdmN72jogE=
-github.com/fluxninja/aperture/api/v2 v2.0.0-20240123210709-ea3fb20237aa h1:eze+1jlQwCkuNPIk4RmTEBY7sLnL/WSDu/vWNLPz43I=
-github.com/fluxninja/aperture/api/v2 v2.0.0-20240123210709-ea3fb20237aa/go.mod h1:KSjIteqXmGJl1WOxQeBF9/K6/0sMHfKsRl5VOQkxyNg=
+github.com/fluxninja/aperture/api/v2 v2.0.0-20240126005715-ea93b670518e h1:KoW+oFm/1S8plY6DvVTJXs0jRgmvtW/psjuLh1yWHZM=
+github.com/fluxninja/aperture/api/v2 v2.0.0-20240126005715-ea93b670518e/go.mod h1:KSjIteqXmGJl1WOxQeBF9/K6/0sMHfKsRl5VOQkxyNg=
github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A=
github.com/go-logr/logr v1.3.0 h1:2y3SDp0ZXuc6/cjLSZ+Q3ir+QB9T/iG5yYRXqsagWSY=
github.com/go-logr/logr v1.3.0/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY=
diff --git a/sdks/aperture-go/go.mod b/sdks/aperture-go/go.mod
index 4f92cc627f..4ce8fec0b4 100644
--- a/sdks/aperture-go/go.mod
+++ b/sdks/aperture-go/go.mod
@@ -3,7 +3,7 @@ module github.com/fluxninja/aperture-go/v2
go 1.21.4
require (
- github.com/fluxninja/aperture/api/v2 v2.0.0-20240123210709-ea3fb20237aa
+ github.com/fluxninja/aperture/api/v2 v2.0.0-20240126005715-ea93b670518e
github.com/gorilla/mux v1.8.1
go.opentelemetry.io/otel v1.21.0
go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.21.0
diff --git a/sdks/aperture-go/go.sum b/sdks/aperture-go/go.sum
index 9c80f945a7..4a627140fc 100644
--- a/sdks/aperture-go/go.sum
+++ b/sdks/aperture-go/go.sum
@@ -4,8 +4,8 @@ github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/envoyproxy/protoc-gen-validate v1.0.2 h1:QkIBuU5k+x7/QXPvPPnWXWlCdaBFApVqftFV6k087DA=
github.com/envoyproxy/protoc-gen-validate v1.0.2/go.mod h1:GpiZQP3dDbg4JouG/NNS7QWXpgx6x8QiMKdmN72jogE=
-github.com/fluxninja/aperture/api/v2 v2.0.0-20240123210709-ea3fb20237aa h1:eze+1jlQwCkuNPIk4RmTEBY7sLnL/WSDu/vWNLPz43I=
-github.com/fluxninja/aperture/api/v2 v2.0.0-20240123210709-ea3fb20237aa/go.mod h1:KSjIteqXmGJl1WOxQeBF9/K6/0sMHfKsRl5VOQkxyNg=
+github.com/fluxninja/aperture/api/v2 v2.0.0-20240126005715-ea93b670518e h1:KoW+oFm/1S8plY6DvVTJXs0jRgmvtW/psjuLh1yWHZM=
+github.com/fluxninja/aperture/api/v2 v2.0.0-20240126005715-ea93b670518e/go.mod h1:KSjIteqXmGJl1WOxQeBF9/K6/0sMHfKsRl5VOQkxyNg=
github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A=
github.com/go-logr/logr v1.3.0 h1:2y3SDp0ZXuc6/cjLSZ+Q3ir+QB9T/iG5yYRXqsagWSY=
github.com/go-logr/logr v1.3.0/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY=
diff --git a/sdks/aperture-js/example/rate_limit_example.ts b/sdks/aperture-js/example/rate_limit_example.ts
index f6dafeb9fe..f7e2da5255 100644
--- a/sdks/aperture-js/example/rate_limit_example.ts
+++ b/sdks/aperture-js/example/rate_limit_example.ts
@@ -31,7 +31,7 @@ async function handleRequestRateLimit(apertureClient: ApertureClient) {
user_id: "some_user_id",
},
grpcCallOptions: {
- deadline: Date.now() + 300, // ms
+ deadline: Date.now() + 1000, // ms
},
});
// END: RLStartFlow
diff --git a/sdks/aperture-py/aperture_sdk/_gen/aperture/flowcontrol/check/v1/check_pb2.py b/sdks/aperture-py/aperture_sdk/_gen/aperture/flowcontrol/check/v1/check_pb2.py
index 190abd7e57..fd04bbe160 100644
--- a/sdks/aperture-py/aperture_sdk/_gen/aperture/flowcontrol/check/v1/check_pb2.py
+++ b/sdks/aperture-py/aperture_sdk/_gen/aperture/flowcontrol/check/v1/check_pb2.py
@@ -1,6 +1,7 @@
# -*- coding: utf-8 -*-
# Generated by the protocol buffer compiler. DO NOT EDIT!
# source: aperture/flowcontrol/check/v1/check.proto
+# Protobuf Python Version: 4.25.2
"""Generated protocol buffer code."""
from google.protobuf import descriptor as _descriptor
from google.protobuf import descriptor_pool as _descriptor_pool
diff --git a/sdks/aperture-py/aperture_sdk/client.py b/sdks/aperture-py/aperture_sdk/client.py
index 0228a5ddab..c277a5d6ba 100644
--- a/sdks/aperture-py/aperture_sdk/client.py
+++ b/sdks/aperture-py/aperture_sdk/client.py
@@ -5,8 +5,7 @@
import logging
import time
import typing
-from dataclasses import dataclass
-from typing import Callable, Dict, Optional, Type, TypeVar
+from typing import Callable, Optional, Type, TypeVar
import grpc
from aperture_sdk._gen.aperture.flowcontrol.check.v1.check_pb2 import (
@@ -16,9 +15,9 @@
from aperture_sdk._gen.aperture.flowcontrol.check.v1.check_pb2_grpc import (
FlowControlServiceStub,
)
+from aperture_sdk.client_common import *
from aperture_sdk.const import (
default_grpc_reconnection_time,
- default_rpc_timeout,
flow_start_timestamp_label,
library_name,
library_version,
@@ -27,7 +26,6 @@
)
from aperture_sdk.flow import Flow
from aperture_sdk.utils import TWrappedReturn, run_fn
-from grpc import AuthMetadataContext, AuthMetadataPluginCallback
from opentelemetry import baggage, trace
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
from opentelemetry.sdk.resources import SERVICE_NAME, SERVICE_VERSION, Resource
@@ -37,26 +35,6 @@
TApertureClient = TypeVar("TApertureClient", bound="ApertureClient")
TWrappedFunction = Callable[..., TWrappedReturn]
-Labels = Dict[str, str]
-
-
-class ApertureCloudAuthMetadataPlugin(grpc.AuthMetadataPlugin):
- def __init__(self, api_key):
- self.api_key = api_key
-
- def __call__(
- self, context: AuthMetadataContext, callback: AuthMetadataPluginCallback
- ) -> None:
- callback((("x-api-key", self.api_key),), None)
-
-
-@dataclass
-class FlowParams:
- explicit_labels: Optional[Labels] = None
- check_timeout: datetime.timedelta = default_rpc_timeout
- ramp_mode: bool = False
- result_cache_key: Optional[str] = None
- global_cache_keys: Optional[typing.List[str]] = None
class ApertureClient:
@@ -69,7 +47,7 @@ def __init__(
channel: grpc.Channel,
otlp_exporter: OTLPSpanExporter,
):
- self.logger = logging.getLogger("aperture-py-sdk")
+ self.logger = logging.getLogger("aperture-py")
resource = Resource.create(
{
@@ -102,7 +80,6 @@ def new_client(
credentials = grpc.ssl_channel_credentials()
if api_key:
metadata_plugin_instance = ApertureCloudAuthMetadataPlugin(api_key)
-
credentials = grpc.composite_channel_credentials(
credentials,
grpc.metadata_call_credentials(
@@ -110,6 +87,7 @@ def new_client(
name="x-api-key",
),
)
+
otlp_exporter = OTLPSpanExporter(
endpoint=address,
insecure=insecure,
@@ -169,11 +147,9 @@ def start_flow(
try:
# stub.Check is typed to accept an int, but it actually accepts a float
timeout = typing.cast(int, params.check_timeout.total_seconds())
- response = (
- stub.Check(request)
- if timeout == 0
- else stub.Check(request, timeout=timeout)
- )
+ if timeout == 0:
+ timeout = None
+ response = stub.Check(request, timeout=timeout)
except grpc.RpcError as e:
self.logger.debug(f"Aperture gRPC call failed: {e.details()}")
response = None
@@ -199,12 +175,12 @@ def decorate(
def decorator(fn: TWrappedFunction) -> TWrappedFunction:
@functools.wraps(fn)
async def wrapper(*args, **kwargs):
- with self.start_flow(control_point, params) as flow:
- if flow.should_run():
- return await run_fn(fn, *args, **kwargs)
- else:
- if on_reject:
- return on_reject()
+ flow = self.start_flow(control_point, params)
+ if flow.should_run():
+ return await run_fn(fn, *args, **kwargs)
+ else:
+ if on_reject:
+ return on_reject()
return wrapper
diff --git a/sdks/aperture-py/aperture_sdk/client_async.py b/sdks/aperture-py/aperture_sdk/client_async.py
new file mode 100644
index 0000000000..b7f8e919cc
--- /dev/null
+++ b/sdks/aperture-py/aperture_sdk/client_async.py
@@ -0,0 +1,190 @@
+"""ApertureClientAsync for starting Flows."""
+
+import datetime
+import functools
+import logging
+import time
+import typing
+from typing import Callable, Optional, Type, TypeVar
+
+import grpc.aio
+from aperture_sdk._gen.aperture.flowcontrol.check.v1.check_pb2 import (
+ CacheLookupRequest,
+ CheckRequest,
+)
+from aperture_sdk._gen.aperture.flowcontrol.check.v1.check_pb2_grpc import (
+ FlowControlServiceStub,
+)
+from aperture_sdk.client_common import *
+from aperture_sdk.const import (
+ default_grpc_reconnection_time,
+ flow_start_timestamp_label,
+ library_name,
+ library_version,
+ source_label,
+ workload_start_timestamp_label,
+)
+from aperture_sdk.flow_async import FlowAsync
+from aperture_sdk.utils import TWrappedReturn, run_fn
+from opentelemetry import baggage, trace
+from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
+from opentelemetry.sdk.resources import SERVICE_NAME, SERVICE_VERSION, Resource
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import BatchSpanProcessor
+from opentelemetry.util import types as otel_types
+
+TApertureClientAsync = TypeVar("TApertureClientAsync", bound="ApertureClientAsync")
+TWrappedFunction = Callable[..., TWrappedReturn]
+
+
+class ApertureClientAsync:
+ """
+ ApertureClientAsync can be used to start Flows.
+ """
+
+ def __init__(
+ self,
+ channel: grpc.aio.Channel,
+ otlp_exporter: OTLPSpanExporter,
+ ):
+ self.logger = logging.getLogger("aperture-py")
+
+ resource = Resource.create(
+ {
+ SERVICE_NAME: library_name,
+ SERVICE_VERSION: library_version,
+ }
+ )
+ tracer_provider = TracerProvider(resource=resource)
+ trace.set_tracer_provider(tracer_provider)
+ self.tracer = trace.get_tracer(library_name, library_version)
+
+ span_processor = BatchSpanProcessor(otlp_exporter)
+ tracer_provider.add_span_processor(span_processor)
+ self.otlp_exporter = otlp_exporter
+ self.grpc_channel = channel
+
+ @classmethod
+ def new_client(
+ cls: Type[TApertureClientAsync],
+ address: str,
+ api_key: Optional[str] = None,
+ insecure: bool = False,
+ grpc_timeout: datetime.timedelta = default_grpc_reconnection_time,
+ credentials: Optional[grpc.ChannelCredentials] = None,
+ compression: grpc.Compression = grpc.Compression.NoCompression,
+ ) -> TApertureClientAsync:
+ if not address:
+ raise ValueError("Address must be provided")
+ if not credentials:
+ credentials = grpc.ssl_channel_credentials()
+ if api_key:
+ metadata_plugin_instance = ApertureCloudAuthMetadataPlugin(api_key)
+ credentials = grpc.composite_channel_credentials(
+ credentials,
+ grpc.metadata_call_credentials(
+ metadata_plugin=metadata_plugin_instance,
+ name="x-api-key",
+ ),
+ )
+
+ otlp_exporter = OTLPSpanExporter(
+ endpoint=address,
+ insecure=insecure,
+ credentials=credentials,
+ compression=compression,
+ timeout=int(grpc_timeout.total_seconds()),
+ )
+ grpc_channel_options_dict = {
+ "grpc.keepalive_time_ms": 10000,
+ "grpc.keepalive_timeout_ms": 5000,
+ }
+ grpc_channel_options = [(k, v) for k, v in grpc_channel_options_dict.items()]
+ grpc_channel = (
+ grpc.aio.insecure_channel(
+ address, compression=compression, options=grpc_channel_options
+ )
+ if insecure
+ else grpc.aio.secure_channel(
+ address,
+ credentials,
+ compression=compression,
+ options=grpc_channel_options,
+ )
+ )
+ return cls(
+ channel=grpc_channel,
+ otlp_exporter=otlp_exporter,
+ )
+
+ async def start_flow(
+ self,
+ control_point: str,
+ params: FlowParams,
+ ) -> FlowAsync:
+ labels: Labels = {}
+ labels.update({key: str(value) for key, value in baggage.get_all().items()})
+ # Explicit labels override baggage
+ labels.update(params.explicit_labels or {})
+ request = CheckRequest(
+ control_point=control_point,
+ labels=labels,
+ ramp_mode=params.ramp_mode,
+ expect_end=True,
+ cache_lookup_request=CacheLookupRequest(
+ result_cache_key=params.result_cache_key,
+ global_cache_keys=params.global_cache_keys,
+ ),
+ )
+ span_attributes: otel_types.Attributes = {
+ flow_start_timestamp_label: time.monotonic_ns(),
+ source_label: "sdk",
+ }
+
+ span = self.tracer.start_span("Aperture Check", attributes=span_attributes)
+ stub = FlowControlServiceStub(self.grpc_channel)
+ error: Optional[Exception] = None
+ try:
+ # stub.Check is typed to accept an int, but it actually accepts a float
+ timeout = typing.cast(int, params.check_timeout.total_seconds())
+ if timeout == 0:
+ timeout = None
+ response = await stub.Check(request, timeout=timeout)
+ except grpc.RpcError as e:
+ self.logger.debug(f"Aperture gRPC call failed: {e.details()}")
+ response = None
+ error = e
+ span.set_attribute(workload_start_timestamp_label, time.monotonic_ns())
+ return FlowAsync(
+ fcs_stub=stub,
+ control_point=control_point,
+ span=span,
+ check_response=response,
+ ramp_mode=params.ramp_mode,
+ cache_key=params.result_cache_key,
+ error=error,
+ grpc_channel=self.grpc_channel,
+ )
+
+ def decorate(
+ self,
+ control_point: str,
+ params: FlowParams = FlowParams(),
+ on_reject: Optional[Callable] = None,
+ ) -> Callable[[TWrappedFunction], TWrappedFunction]:
+ def decorator(fn: TWrappedFunction) -> TWrappedFunction:
+ @functools.wraps(fn)
+ async def wrapper(*args, **kwargs):
+ flow = await self.start_flow(control_point, params)
+ if flow.should_run():
+ return await run_fn(fn, *args, **kwargs)
+ else:
+ if on_reject:
+ return on_reject()
+
+ return wrapper
+
+ return decorator
+
+ def close(self):
+ self.otlp_exporter.shutdown()
diff --git a/sdks/aperture-py/aperture_sdk/client_common.py b/sdks/aperture-py/aperture_sdk/client_common.py
new file mode 100644
index 0000000000..82f095cdbf
--- /dev/null
+++ b/sdks/aperture-py/aperture_sdk/client_common.py
@@ -0,0 +1,27 @@
+import datetime
+from dataclasses import dataclass
+from typing import Dict, List, Optional
+
+from aperture_sdk.const import default_rpc_timeout
+from grpc import AuthMetadataContext, AuthMetadataPlugin, AuthMetadataPluginCallback
+
+Labels = Dict[str, str]
+
+
+class ApertureCloudAuthMetadataPlugin(AuthMetadataPlugin):
+ def __init__(self, api_key):
+ self.api_key = api_key
+
+ def __call__(
+ self, context: AuthMetadataContext, callback: AuthMetadataPluginCallback
+ ) -> None:
+ callback((("x-api-key", self.api_key),), None)
+
+
+@dataclass
+class FlowParams:
+ explicit_labels: Optional[Labels] = None
+ check_timeout: datetime.timedelta = default_rpc_timeout
+ ramp_mode: bool = False
+ result_cache_key: Optional[str] = None
+ global_cache_keys: Optional[List[str]] = None
diff --git a/sdks/aperture-py/aperture_sdk/flow.py b/sdks/aperture-py/aperture_sdk/flow.py
index 414b783b18..d74f4e9388 100644
--- a/sdks/aperture-py/aperture_sdk/flow.py
+++ b/sdks/aperture-py/aperture_sdk/flow.py
@@ -14,7 +14,6 @@
CacheEntry,
CacheUpsertRequest,
FlowEndRequest,
- FlowEndResponse,
InflightRequestRef,
StatusCode,
)
@@ -27,36 +26,11 @@
flow_end_timestamp_label,
flow_status_label,
)
+from aperture_sdk.flow_common import *
from google.protobuf import json_format
from google.protobuf.duration_pb2 import Duration
from opentelemetry import trace
-
-class FlowDecision(enum.Enum):
- Accepted = enum.auto()
- Rejected = enum.auto()
- Unreachable = enum.auto()
-
-
-class FlowStatus(enum.Enum):
- OK = enum.auto()
- Error = enum.auto()
-
-
-class EndResponse:
- def __init__(
- self, error: Optional[Exception], flow_end_response: Optional[FlowEndResponse]
- ):
- self.error = error
- self.flow_end_response = flow_end_response
-
- def get_error(self) -> Optional[Exception]:
- return self.error
-
- def get_flow_end_response(self) -> Optional[FlowEndResponse]:
- return self.flow_end_response
-
-
TFlow = TypeVar("TFlow", bound="Flow")
@@ -160,7 +134,7 @@ def end(self) -> EndResponse:
if self.check_response:
for decision in self.check_response.limiter_decisions:
- if decision.concurrency_limiter_info:
+ if decision.WhichOneof("details") == "concurrency_limiter_info":
ref: InflightRequestRef = InflightRequestRef(
policy_name=decision.policy_name,
policy_hash=decision.policy_hash,
@@ -168,14 +142,12 @@ def end(self) -> EndResponse:
label=decision.concurrency_limiter_info.label,
request_id=decision.concurrency_limiter_info.request_id,
)
-
if decision.concurrency_limiter_info.tokens_info:
ref.tokens = (
decision.concurrency_limiter_info.tokens_info.consumed
)
inflight_request_ref.append(ref)
-
- if decision.concurrency_scheduler_info:
+ elif decision.WhichOneof("details") == "concurrency_scheduler_info":
ref: InflightRequestRef = InflightRequestRef(
policy_name=decision.policy_name,
policy_hash=decision.policy_hash,
@@ -183,7 +155,6 @@ def end(self) -> EndResponse:
label=decision.concurrency_scheduler_info.label,
request_id=decision.concurrency_scheduler_info.request_id,
)
-
if decision.concurrency_scheduler_info.tokens_info:
ref.tokens = (
decision.concurrency_scheduler_info.tokens_info.consumed
diff --git a/sdks/aperture-py/aperture_sdk/flow_async.py b/sdks/aperture-py/aperture_sdk/flow_async.py
new file mode 100644
index 0000000000..96bbc17505
--- /dev/null
+++ b/sdks/aperture-py/aperture_sdk/flow_async.py
@@ -0,0 +1,389 @@
+"""FlowAsync started using the ApertureClientAsync."""
+
+import datetime
+import logging
+import time
+from contextlib import AbstractContextManager
+from typing import List, Optional, TypeVar
+
+import grpc.aio
+from aperture_sdk._gen.aperture.flowcontrol.check.v1 import check_pb2
+from aperture_sdk._gen.aperture.flowcontrol.check.v1.check_pb2 import (
+ CacheDeleteRequest,
+ CacheDeleteResponse,
+ CacheEntry,
+ CacheUpsertRequest,
+ FlowEndRequest,
+ InflightRequestRef,
+ StatusCode,
+)
+from aperture_sdk._gen.aperture.flowcontrol.check.v1.check_pb2_grpc import (
+ FlowControlServiceStub,
+)
+from aperture_sdk.cache import *
+from aperture_sdk.const import (
+ check_response_label,
+ flow_end_timestamp_label,
+ flow_status_label,
+)
+from aperture_sdk.flow_common import *
+from google.protobuf import json_format
+from google.protobuf.duration_pb2 import Duration
+from opentelemetry import trace
+
+TFlowAsync = TypeVar("TFlowAsync", bound="FlowAsync")
+
+
+class FlowAsync(AbstractContextManager):
+ def __init__(
+ self,
+ fcs_stub: FlowControlServiceStub,
+ control_point: str,
+ span: trace.Span,
+ check_response: Optional[check_pb2.CheckResponse],
+ ramp_mode: bool,
+ cache_key: Optional[str],
+ error: Optional[Exception],
+ grpc_channel: grpc.aio.Channel,
+ ):
+ self._fcs_stub = fcs_stub
+ self._control_point = control_point
+ self._span = span
+ self._check_response = check_response
+ self._cache_key = cache_key
+ self._status_code = FlowStatus.OK
+ self._ended = False
+ self._ramp_mode = ramp_mode
+ self._error = error
+ self._grpc_channel = grpc_channel
+ self.logger = logging.getLogger("aperture-py-sdk-flow")
+
+ def should_run(self) -> bool:
+ return self.decision == FlowDecision.Accepted or (
+ (not self._ramp_mode) and self.decision == FlowDecision.Unreachable
+ )
+
+ def http_response_code(self) -> Optional[int]:
+ if self.check_response is None:
+ return 200 # Default to 200 if check_response is None
+
+ if self.check_response.denied_response_status_code == StatusCode.Empty:
+ return 200 # Default to 200 if denied_response_status_code is None or Empty
+
+ return self.check_response.denied_response_status_code
+
+ def retry_after(self) -> Optional[datetime.timedelta]:
+ if self.check_response is None:
+ return None
+ if self.check_response.wait_time is None:
+ return None
+ return datetime.timedelta(seconds=self.check_response.wait_time.seconds)
+
+ @property
+ def decision(self) -> FlowDecision:
+ if self.check_response is None:
+ return FlowDecision.Unreachable
+ if (
+ self.check_response.decision_type
+ == check_pb2.CheckResponse.DECISION_TYPE_ACCEPTED
+ ):
+ return FlowDecision.Accepted
+ return FlowDecision.Rejected
+
+ @property
+ def success(self) -> bool:
+ return self.decision != FlowDecision.Unreachable
+
+ @property
+ def check_response(self) -> Optional[check_pb2.CheckResponse]:
+ return self._check_response
+
+ def set_status(self, status_code: FlowStatus) -> None:
+ self._status_code = status_code
+
+ async def end(self) -> EndResponse:
+ if self._ended:
+ self.logger.warning("attempting to end an already ended flow")
+ return EndResponse(
+ error=Exception("attempting to end an already ended flow"),
+ flow_end_response=None,
+ )
+ if not self.check_response:
+ self.logger.warning("attempting to end a flow with no check response")
+ return EndResponse(
+ error=Exception("attempting to end a flow with no check response"),
+ flow_end_response=None,
+ )
+ self._ended = True
+
+ check_response_json = (
+ json_format.MessageToJson(self.check_response)
+ if self.check_response
+ else ""
+ )
+ self._span.set_attributes(
+ {
+ flow_status_label: self._status_code.name,
+ check_response_label: check_response_json,
+ flow_end_timestamp_label: time.monotonic_ns(),
+ }
+ )
+ self._span.end()
+
+ inflight_request_ref: List[InflightRequestRef] = list()
+
+ if self.check_response:
+ for decision in self.check_response.limiter_decisions:
+ if decision.WhichOneof("details") == "concurrency_limiter_info":
+ if decision.concurrency_limiter_info.request_id == "":
+ continue
+
+ ref: InflightRequestRef = InflightRequestRef(
+ policy_name=decision.policy_name,
+ policy_hash=decision.policy_hash,
+ component_id=decision.component_id,
+ label=decision.concurrency_limiter_info.label,
+ request_id=decision.concurrency_limiter_info.request_id,
+ )
+ if decision.concurrency_limiter_info.tokens_info:
+ ref.tokens = (
+ decision.concurrency_limiter_info.tokens_info.consumed
+ )
+ inflight_request_ref.append(ref)
+ elif decision.WhichOneof("details") == "concurrency_scheduler_info":
+ if decision.concurrency_scheduler_info.request_id == "":
+ continue
+
+ ref: InflightRequestRef = InflightRequestRef(
+ policy_name=decision.policy_name,
+ policy_hash=decision.policy_hash,
+ component_id=decision.component_id,
+ label=decision.concurrency_scheduler_info.label,
+ request_id=decision.concurrency_scheduler_info.request_id,
+ )
+ if decision.concurrency_scheduler_info.tokens_info:
+ ref.tokens = (
+ decision.concurrency_scheduler_info.tokens_info.consumed
+ )
+ inflight_request_ref.append(ref)
+
+ if inflight_request_ref:
+ flow_end_request = FlowEndRequest(
+ control_point=self._control_point,
+ inflight_requests=inflight_request_ref,
+ )
+
+ try:
+ res = await self._fcs_stub.FlowEnd(flow_end_request)
+ except grpc.RpcError as e:
+ self.logger.error(f"Aperture gRPC call failed: {e.details()}")
+ return EndResponse(
+ error=e,
+ flow_end_response=None,
+ )
+
+ # print("Check response ", self.check_response)
+
+ return EndResponse(
+ error=None,
+ flow_end_response=res,
+ )
+ else:
+ return EndResponse(
+ error=None,
+ flow_end_response=None,
+ )
+
+ def error(self) -> Optional[Exception]:
+ return self._error
+
+ async def set_result_cache(
+ self, value: str, ttl: datetime.timedelta, **grpc_opts
+ ) -> KeyUpsertResponse:
+ if not self._cache_key:
+ return KeyUpsertResponse(ValueError("No cache key"))
+
+ ttl_duration = Duration()
+ ttl_duration.FromTimedelta(ttl)
+ cache_upsert_request = CacheUpsertRequest(
+ control_point=self._control_point,
+ result_cache_entry=CacheEntry(
+ key=self._cache_key,
+ value=bytes(value, "utf-8"),
+ ttl=ttl_duration,
+ ),
+ )
+
+ try:
+ res = await self._fcs_stub.CacheUpsert(cache_upsert_request, **grpc_opts)
+ except grpc.RpcError as e:
+ self.logger.debug(f"Aperture gRPC call failed: {e.details()}")
+ return KeyUpsertResponse(e)
+
+ if res.result_cache_response is None:
+ return KeyUpsertResponse(ValueError("No cache upsert response"))
+
+ return KeyUpsertResponse(
+ convert_cache_error(res.result_cache_response.error),
+ )
+
+ async def delete_result_cache(self, **grpc_opts) -> KeyDeleteResponse:
+ if not self._cache_key:
+ return KeyDeleteResponse(ValueError("No cache key"))
+
+ cache_delete_request = CacheDeleteRequest(
+ control_point=self._control_point,
+ result_cache_key=self._cache_key,
+ )
+
+ try:
+ res: CacheDeleteResponse = await self._fcs_stub.CacheDelete(
+ cache_delete_request, **grpc_opts
+ )
+ except grpc.RpcError as e:
+ self.logger.debug(f"Aperture gRPC call failed: {e.details()}")
+ return KeyDeleteResponse(e)
+
+ if res.result_cache_response is None:
+ return KeyDeleteResponse(ValueError("No cache delete response"))
+
+ return KeyDeleteResponse(
+ convert_cache_error(res.result_cache_response.error),
+ )
+
+ def result_cache(self) -> KeyLookupResponse:
+ if self._error is not None:
+ return KeyLookupResponse(None, LookupStatus.MISS, self._error)
+ if not self.check_response:
+ return KeyLookupResponse(
+ None, LookupStatus.MISS, ValueError("response is null")
+ )
+ if not self.should_run():
+ return KeyLookupResponse(
+ None, LookupStatus.MISS, ValueError("flow was rejected")
+ )
+
+ if (
+ not self.check_response.cache_lookup_response
+ or not self.check_response.cache_lookup_response.result_cache_response
+ ):
+ return KeyLookupResponse(
+ None, LookupStatus.MISS, ValueError("No cache lookup response")
+ )
+ lookup_response = (
+ self.check_response.cache_lookup_response.result_cache_response
+ )
+ return KeyLookupResponse(
+ lookup_response.value,
+ convert_cache_lookup_status(lookup_response.lookup_status),
+ convert_cache_error(lookup_response.error),
+ )
+
+ async def set_global_cache(
+ self, key: str, value: str, ttl: datetime.timedelta, **grpc_opts
+ ) -> KeyUpsertResponse:
+ ttl_duration = Duration()
+ ttl_duration.FromTimedelta(ttl)
+ cache_upsert_request = CacheUpsertRequest(
+ global_cache_entries={
+ key: CacheEntry(
+ value=bytes(value, "utf-8"),
+ ttl=ttl_duration,
+ ),
+ },
+ )
+
+ try:
+ res = await self._fcs_stub.CacheUpsert(cache_upsert_request, **grpc_opts)
+ except grpc.RpcError as e:
+ self.logger.debug(f"Aperture gRPC call failed: {e.details()}")
+ return KeyUpsertResponse(e)
+
+ responses = res.global_cache_responses
+ if responses is None:
+ return KeyUpsertResponse(ValueError("No cache upsert response"))
+ if key not in responses:
+ return KeyUpsertResponse(
+ ValueError("Key missing from global cache response")
+ )
+
+ return KeyUpsertResponse(
+ convert_cache_error(responses[key].error),
+ )
+
+ async def delete_global_cache(self, key: str, **grpc_opts) -> KeyDeleteResponse:
+ cache_delete_request = CacheDeleteRequest(
+ global_cache_keys=[key],
+ )
+
+ try:
+ res: CacheDeleteResponse = await self._fcs_stub.CacheDelete(
+ cache_delete_request, **grpc_opts
+ )
+ except grpc.RpcError as e:
+ self.logger.debug(f"Aperture gRPC call failed: {e.details()}")
+ return KeyDeleteResponse(e)
+
+ delete_responses = res.global_cache_responses
+
+ if delete_responses is None:
+ return KeyDeleteResponse(ValueError("No cache delete response"))
+ if key not in delete_responses:
+ return KeyDeleteResponse(
+ ValueError("Key missing from global cache response")
+ )
+
+ return KeyDeleteResponse(
+ convert_cache_error(delete_responses[key].error),
+ )
+
+ def global_cache(self, key: str) -> KeyLookupResponse:
+ if self._error is not None:
+ return KeyLookupResponse(None, LookupStatus.MISS, self._error)
+ if not self.check_response:
+ return KeyLookupResponse(
+ None, LookupStatus.MISS, ValueError("response is null")
+ )
+ if not self.should_run():
+ return KeyLookupResponse(
+ None, LookupStatus.MISS, ValueError("flow was rejected")
+ )
+
+ if (
+ not self.check_response.cache_lookup_response
+ or not self.check_response.cache_lookup_response.global_cache_responses
+ ):
+ return KeyLookupResponse(
+ None, LookupStatus.MISS, ValueError("No global cache lookup response")
+ )
+
+ lookup_response_map = (
+ self.check_response.cache_lookup_response.global_cache_responses
+ )
+ if key not in lookup_response_map:
+ return KeyLookupResponse(
+ None, LookupStatus.MISS, ValueError("Unknown global cache key")
+ )
+
+ lookup_response = lookup_response_map[key]
+ return KeyLookupResponse(
+ lookup_response.value,
+ convert_cache_lookup_status(lookup_response.lookup_status),
+ convert_cache_error(lookup_response.error),
+ )
+
+ def __enter__(self: TFlowAsync) -> TFlowAsync:
+ return self
+
+ async def __exit__(self, exc_type, _exc_value, _traceback) -> None:
+ if self._ended:
+ return
+ if exc_type is not None:
+ self.set_status(FlowStatus.Error)
+ res = await self.end()
+
+ if res.get_error():
+ self.logger.warning(f"Failed to end flow: {res.get_error()}")
+
+ if res.get_flow_end_response():
+ print(f"Ended flow: {res.get_flow_end_response()}")
diff --git a/sdks/aperture-py/aperture_sdk/flow_common.py b/sdks/aperture-py/aperture_sdk/flow_common.py
new file mode 100644
index 0000000000..b40b735509
--- /dev/null
+++ b/sdks/aperture-py/aperture_sdk/flow_common.py
@@ -0,0 +1,29 @@
+import enum
+from typing import Optional
+
+from aperture_sdk._gen.aperture.flowcontrol.check.v1.check_pb2 import FlowEndResponse
+
+
+class FlowDecision(enum.Enum):
+ Accepted = enum.auto()
+ Rejected = enum.auto()
+ Unreachable = enum.auto()
+
+
+class FlowStatus(enum.Enum):
+ OK = enum.auto()
+ Error = enum.auto()
+
+
+class EndResponse:
+ def __init__(
+ self, error: Optional[Exception], flow_end_response: Optional[FlowEndResponse]
+ ):
+ self.error = error
+ self.flow_end_response = flow_end_response
+
+ def get_error(self) -> Optional[Exception]:
+ return self.error
+
+ def get_flow_end_response(self) -> Optional[FlowEndResponse]:
+ return self.flow_end_response
diff --git a/sdks/aperture-py/docs/all-documents.html b/sdks/aperture-py/docs/all-documents.html
index 84867e32b9..157c71ae66 100644
--- a/sdks/aperture-py/docs/all-documents.html
+++ b/sdks/aperture-py/docs/all-documents.html
@@ -27,7 +27,7 @@