forked from feast-dev/feast
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathfeast_core_server.py
More file actions
126 lines (99 loc) · 4.12 KB
/
Copy pathfeast_core_server.py
File metadata and controls
126 lines (99 loc) · 4.12 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
import logging
import time
from concurrent import futures
import grpc
from google.protobuf.timestamp_pb2 import Timestamp
import feast.core.CoreService_pb2_grpc as Core
from feast.core.CoreService_pb2 import (
ApplyFeatureSetRequest,
ApplyFeatureSetResponse,
GetFeastCoreVersionResponse,
ListFeatureSetsRequest,
ListFeatureSetsResponse,
)
from feast.core.FeatureSet_pb2 import FeatureSet as FeatureSetProto
from feast.core.FeatureSet_pb2 import FeatureSetMeta, FeatureSetStatus
from feast.core.Source_pb2 import KafkaSourceConfig as KafkaSourceConfigProto
from feast.core.Source_pb2 import SourceType as SourceTypeProto
_logger = logging.getLogger(__name__)
_ONE_DAY_IN_SECONDS = 60 * 60 * 24
_SIGNATURE_HEADER_KEY = "authorization"
class DisallowAuthInterceptor(grpc.ServerInterceptor):
def __init__(self):
def abort(ignored_request, context):
context.abort(grpc.StatusCode.UNAUTHENTICATED, "Invalid signature")
self._abortion = grpc.unary_unary_rpc_method_handler(abort)
def intercept_service(self, continuation, handler_call_details):
print(handler_call_details.invocation_metadata)
if "Bearer" in handler_call_details.invocation_metadata[0][1]:
return self._abortion
else:
return continuation(handler_call_details)
class AllowAuthInterceptor(grpc.ServerInterceptor):
def __init__(self):
def abort(ignored_request, context):
context.abort(grpc.StatusCode.UNAUTHENTICATED, "Invalid signature")
self._abortion = grpc.unary_unary_rpc_method_handler(abort)
def intercept_service(self, continuation, handler_call_details):
print(handler_call_details.invocation_metadata)
if "Bearer" in handler_call_details.invocation_metadata[0][1]:
return continuation(handler_call_details)
else:
return self._abortion
class CoreServicer(Core.CoreServiceServicer):
def __init__(self):
self._feature_sets = dict()
def GetFeastCoreVersion(self, request, context):
return GetFeastCoreVersionResponse(version="0.3.2")
def ListFeatureSets(self, request: ListFeatureSetsRequest, context):
filtered_feature_set_response = [
fs
for fs in list(self._feature_sets.values())
if (
not request.filter.feature_set_name
or request.filter.feature_set_name == "*"
or fs.spec.name == request.filter.feature_set_name
)
]
return ListFeatureSetsResponse(feature_sets=filtered_feature_set_response)
def ApplyFeatureSet(self, request: ApplyFeatureSetRequest, context):
feature_set = request.feature_set
if feature_set.spec.source.type == SourceTypeProto.INVALID:
feature_set.spec.source.kafka_source_config.CopyFrom(
KafkaSourceConfigProto(bootstrap_servers="server.com", topic="topic1")
)
feature_set.spec.source.type = SourceTypeProto.KAFKA
feature_set_meta = FeatureSetMeta(
status=FeatureSetStatus.STATUS_READY,
created_timestamp=Timestamp(seconds=10),
)
applied_feature_set = FeatureSetProto(
spec=feature_set.spec, meta=feature_set_meta
)
self._feature_sets[feature_set.spec.name] = applied_feature_set
_logger.info(
"registered feature set "
+ feature_set.spec.name
+ " with "
+ str(len(feature_set.spec.entities))
+ " entities and "
+ str(len(feature_set.spec.features))
+ " features"
)
return ApplyFeatureSetResponse(
feature_set=applied_feature_set,
status=ApplyFeatureSetResponse.Status.CREATED,
)
def serve():
server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
Core.add_CoreServiceServicer_to_server(CoreServicer(), server)
server.add_insecure_port("[::]:50051")
server.start()
try:
while True:
time.sleep(_ONE_DAY_IN_SECONDS)
except KeyboardInterrupt:
server.stop(0)
if __name__ == "__main__":
logging.basicConfig()
serve()