nucliadb-utils 6.2.1.post3165__py3-none-any.whl → 6.2.1.post3188__py3-none-any.whl

nucliadb_utils/audit/stream.py

@@ -37,15 +37,17 @@ from starlette.types import ASGIApp
 
 from nucliadb_protos.audit_pb2 import AuditField, AuditRequest, ChatContext, ClientType, RetrievedContext
 from nucliadb_protos.kb_usage_pb2 import (
-    ClientType as ClientTypeKbUsage,
-)
-from nucliadb_protos.kb_usage_pb2 import (
+    ActivityLogMatch,
+    ActivityLogMatchType,
     KBSource,
     Search,
     SearchType,
     Service,
     Storage,
 )
+from nucliadb_protos.kb_usage_pb2 import (
+    ClientType as ClientTypeKbUsage,
+)
 from nucliadb_protos.nodereader_pb2 import SearchRequest
 from nucliadb_protos.resources_pb2 import FieldID
 from nucliadb_utils import logger
@@ -64,10 +66,10 @@ class RequestContext:
 request_context_var = contextvars.ContextVar[Optional[RequestContext]]("request_context", default=None)
 
 
-def get_trace_id() -> str:
+def get_trace_id() -> Optional[str]:
     span = get_current_span()
     if span is None:
-        return ""
+        return None
     return format_trace_id(span.get_span_context().trace_id)
 
 
@@ -88,7 +90,7 @@ class AuditMiddleware(BaseHTTPMiddleware):
         context = RequestContext()
         token = request_context_var.set(context)
         context.audit_request.time.FromDatetime(datetime.now(tz=timezone.utc))
-        context.audit_request.trace_id = get_trace_id()
+        context.audit_request.trace_id = get_trace_id() or ""
         context.path = request.url.path
 
         if request.url.path.split("/")[-1] in ("ask", "search", "find"):
@@ -252,7 +254,7 @@ class StreamAuditStorage(AuditStorage):
 
         # Reports MODIFIED / DELETED / NEW events
 
-        auditrequest.trace_id = get_trace_id()
+        auditrequest.trace_id = get_trace_id() or ""
         auditrequest.kbid = kbid
         auditrequest.userid = user or ""
         auditrequest.rid = rid or ""
@@ -270,12 +272,16 @@ class StreamAuditStorage(AuditStorage):
         self.send(auditrequest)
 
     def report_storage(self, kbid: str, paragraphs: int, fields: int, bytes: int):
+        trace_id = get_trace_id()
         self.kb_usage_utility.send_kb_usage(
             service=Service.NUCLIA_DB,
             account_id=None,
             kb_id=kbid,
             kb_source=KBSource.HOSTED,
             storage=Storage(paragraphs=paragraphs, fields=fields, bytes=bytes),
+            activity_log_match=ActivityLogMatch(id=trace_id, type=ActivityLogMatchType.TRACE_ID)
+            if trace_id
+            else None,
         )
 
     def report_resources(
@@ -284,12 +290,16 @@ class StreamAuditStorage(AuditStorage):
         kbid: str,
         resources: int,
     ):
+        trace_id = get_trace_id()
         self.kb_usage_utility.send_kb_usage(
             service=Service.NUCLIA_DB,
             account_id=None,
             kb_id=kbid,
             kb_source=KBSource.HOSTED,
             storage=Storage(resources=resources),
+            activity_log_match=ActivityLogMatch(id=trace_id, type=ActivityLogMatchType.TRACE_ID)
+            if trace_id
+            else None,
         )
 
     def visited(
@@ -312,12 +322,16 @@ class StreamAuditStorage(AuditStorage):
         auditrequest.type = AuditRequest.VISITED
 
     def delete_kb(self, kbid: str):
+        trace_id = get_trace_id()
         self.kb_usage_utility.send_kb_usage(
             service=Service.NUCLIA_DB,
             account_id=None,
             kb_id=kbid,
             kb_source=KBSource.HOSTED,
             storage=Storage(paragraphs=0, fields=0, resources=0),
+            activity_log_match=ActivityLogMatch(id=trace_id, type=ActivityLogMatchType.TRACE_ID)
+            if trace_id
+            else None,
         )
 
     def search(
@@ -348,6 +362,7 @@ class StreamAuditStorage(AuditStorage):
         else:
             auditrequest.type = AuditRequest.SEARCH
 
+        trace_id = get_trace_id()
         self.kb_usage_utility.send_kb_usage(
             service=Service.NUCLIA_DB,
             account_id=None,
@@ -362,6 +377,9 @@ class StreamAuditStorage(AuditStorage):
                     num_searches=1,
                 )
             ],
+            activity_log_match=ActivityLogMatch(id=trace_id, type=ActivityLogMatchType.TRACE_ID)
+            if trace_id
+            else None,
         )
 
     def chat(

nucliadb_utils/authentication.py

@@ -32,8 +32,9 @@ from starlette.websockets import WebSocket
 
 
 class NucliaUser(BaseUser):
-    def __init__(self, username: str) -> None:
+    def __init__(self, username: str, security_groups: Optional[list[str]] = None) -> None:
         self.username = username
+        self._security_groups = security_groups
 
     @property
     def is_authenticated(self) -> bool:
@@ -43,6 +44,10 @@ class NucliaUser(BaseUser):
     def display_name(self) -> str:
         return self.username
 
+    @property
+    def security_groups(self) -> Optional[list[str]]:
+        return self._security_groups
+
 
 STFUser = NucliaUser
 
@@ -52,9 +57,11 @@ class NucliaCloudAuthenticationBackend(AuthenticationBackend):
         self,
         roles_header: str = "X-NUCLIADB-ROLES",
         user_header: str = "X-NUCLIADB-USER",
+        security_groups_header: str = "X-NUCLIADB-SECURITY-GROUPS",
     ) -> None:
         self.roles_header = roles_header
         self.user_header = user_header
+        self.security_groups_header = security_groups_header
 
     async def authenticate(self, request: HTTPConnection) -> Optional[Tuple[AuthCredentials, BaseUser]]:
         if self.roles_header not in request.headers:
@@ -66,7 +73,15 @@ class NucliaCloudAuthenticationBackend(AuthenticationBackend):
 
         if self.user_header in request.headers:
             user = request.headers[self.user_header]
-            nuclia_user: BaseUser = NucliaUser(username=user)
+
+            raw_security_groups: Optional[str] = request.headers.get(self.security_groups_header)
+
+            security_groups: Optional[list[str]] = None
+            if raw_security_groups is not None:
+                security_groups = raw_security_groups.split(";")
+
+            nuclia_user: NucliaUser = NucliaUser(username=user, security_groups=security_groups)
+
         else:
             nuclia_user = NucliaUser(username="anonymous")
 

nucliadb_utils/nuclia_usage/utils/kb_usage_report.py

@@ -27,6 +27,7 @@ from typing import Optional
 from nats.js.client import JetStreamContext
 
 from nucliadb_protos.kb_usage_pb2 import (
+    ActivityLogMatch,
     KBSource,
     KbUsage,
     Predict,
@@ -96,6 +97,7 @@ class KbUsageReportUtility:
         predicts: Iterable[Predict] = (),
         searches: Iterable[Search] = (),
         storage: Optional[Storage] = None,
+        activity_log_match: Optional[ActivityLogMatch] = None,
     ):
         usage = KbUsage()
         usage.service = service  # type: ignore
@@ -105,6 +107,10 @@ class KbUsageReportUtility:
         if kb_id is not None:
             usage.kb_id = kb_id
         usage.kb_source = kb_source  # type: ignore
+        if activity_log_match is not None:
+            if activity_log_match.id:
+                usage.activity_log_match.id = activity_log_match.id
+                usage.activity_log_match.type = activity_log_match.type
 
         usage.processes.extend(processes)
         usage.predicts.extend(predicts)

{nucliadb_utils-6.2.1.post3165.dist-info → nucliadb_utils-6.2.1.post3188.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: nucliadb_utils
-Version: 6.2.1.post3165
+Version: 6.2.1.post3188
 Home-page: https://nuclia.com
 License: BSD
 Classifier: Development Status :: 4 - Beta
@@ -23,8 +23,8 @@ Requires-Dist: nats-py[nkeys]>=2.6.0
 Requires-Dist: PyNaCl
 Requires-Dist: pyjwt>=2.4.0
 Requires-Dist: mrflagly>=0.2.9
-Requires-Dist: nucliadb-protos>=6.2.1.post3165
-Requires-Dist: nucliadb-telemetry>=6.2.1.post3165
+Requires-Dist: nucliadb-protos>=6.2.1.post3188
+Requires-Dist: nucliadb-telemetry>=6.2.1.post3188
 Provides-Extra: cache
 Requires-Dist: redis>=4.3.4; extra == "cache"
 Requires-Dist: orjson>=3.6.7; extra == "cache"

{nucliadb_utils-6.2.1.post3165.dist-info → nucliadb_utils-6.2.1.post3188.dist-info}/RECORD

@@ -1,6 +1,6 @@
 nucliadb_utils/__init__.py,sha256=EvBCH1iTODe-AgXm48aj4kVUt_Std3PeL8QnwimR5wI,895
 nucliadb_utils/asyncio_utils.py,sha256=h8Y-xpcFFRgNzaiIW0eidz7griAQa7ggbNk34-tAt2c,2888
-nucliadb_utils/authentication.py,sha256=N__d2Ez3JHJv5asYK5TgUcIkKqcAC8ZTLlnfLhfSneM,5837
+nucliadb_utils/authentication.py,sha256=5_b323v2ylJaJvM_0coeSQEtnD-p9IGD-6CPA6IXhik,6471
 nucliadb_utils/const.py,sha256=ai2YJIv0uVOOowhbtSOqcrGMV6UJPjhGhlIRr7ZT8bg,2530
 nucliadb_utils/debug.py,sha256=Q56Nx9Dp7V2ae3CU2H0ztaZcHTJXdlflPLKLeOPZ170,2436
 nucliadb_utils/exceptions.py,sha256=y_3wk77WLVUtdo-5FtbBsdSkCtK_DsJkdWb5BoPn3qo,1094
@@ -24,7 +24,7 @@ nucliadb_utils/aiopynecone/models.py,sha256=XkNIZx4bxdbVo9zYVn8IRp70q4DWUMWN79yb
 nucliadb_utils/audit/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
 nucliadb_utils/audit/audit.py,sha256=VDUqcSToEoz5m7iG-4omjTlsBSEJGcp0_xl51idnU3Q,3384
 nucliadb_utils/audit/basic.py,sha256=pVf1oKp7T_oTDCyt8QbSsfC_uWjF66ooF6rqY-SWSrc,4032
-nucliadb_utils/audit/stream.py,sha256=zT--ECfFNfhL22hhN4IKkPrCaNITKijmufXVoVWm_4c,15233
+nucliadb_utils/audit/stream.py,sha256=AT4OuLTtQ3mID8Geasn6d7qRK8D7MIMUi7GemzU_7mk,16017
 nucliadb_utils/cache/__init__.py,sha256=itSI7dtTwFP55YMX4iK7JzdMHS5CQVUiB1XzQu4UBh8,833
 nucliadb_utils/cache/exceptions.py,sha256=Zu-O_-0-yctOEgoDGI92gPzWfBMRrpiAyESA62ld6MA,975
 nucliadb_utils/cache/nats.py,sha256=-AjCfkFgKVdJUlGR0hT9JDSNkPVFg4S6w9eW-ZIcXPM,7037
@@ -38,7 +38,7 @@ nucliadb_utils/fastapi/run.py,sha256=OVuzjfsX_8IdVXKuK-RsIVf2OdP6t8h9DLfEn2P_-KI
 nucliadb_utils/fastapi/versioning.py,sha256=pwiwuesJW1jElUUI3pI5kcxkigfGBvI64IL6QCBEWS8,3805
 nucliadb_utils/nuclia_usage/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
 nucliadb_utils/nuclia_usage/utils/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
-nucliadb_utils/nuclia_usage/utils/kb_usage_report.py,sha256=F6dhQWhrWzl3qfAdvrqV08hnFKcvFlU8IWN2lzYX0vg,3628
+nucliadb_utils/nuclia_usage/utils/kb_usage_report.py,sha256=6lLuxCCPQVn3dOuZNL5ThPjl2yws-1TJ_7duhQSWkPU,3934
 nucliadb_utils/storages/__init__.py,sha256=5Qc8AUWiJv9_JbGCBpAn88AIJhwDlm0OPQpg2ZdRL4U,872
 nucliadb_utils/storages/azure.py,sha256=FK4c_v9AUAwagScm_F1uDmJeQQq7P4jZswiD2trwb4A,16394
 nucliadb_utils/storages/exceptions.py,sha256=_xUDbfr3XkmUc0xXYZu3DpG1C8TtfmEweRwj9mbRvMU,2507
@@ -59,8 +59,8 @@ nucliadb_utils/tests/indexing.py,sha256=YW2QhkhO9Q_8A4kKWJaWSvXvyQ_AiAwY1VylcfVQ
 nucliadb_utils/tests/local.py,sha256=fXIBasrvdaFJM-sw2wk1_oiFzBcm9O10iCyC-OiXwY8,1914
 nucliadb_utils/tests/nats.py,sha256=RWHjwqq5esuO7OFbP24yYX1cXnpPLcWJwDUdmwCpH28,1897
 nucliadb_utils/tests/s3.py,sha256=pl-RJFjA4MH6iXkqhsh5g8gDuEhrYu1nPZ-laxlrMlE,3704
-nucliadb_utils-6.2.1.post3165.dist-info/METADATA,sha256=ygiFfmhuQ9oD0bhv3ZAtC2In3ZZEORSKUe3U88Z286g,2172
-nucliadb_utils-6.2.1.post3165.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-nucliadb_utils-6.2.1.post3165.dist-info/top_level.txt,sha256=fE3vJtALTfgh7bcAWcNhcfXkNPp_eVVpbKK-2IYua3E,15
-nucliadb_utils-6.2.1.post3165.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-nucliadb_utils-6.2.1.post3165.dist-info/RECORD,,
+nucliadb_utils-6.2.1.post3188.dist-info/METADATA,sha256=Nio-K2aVS2RnjP-9R3mCgisc4-uZBqb5Jf1AvRlbzgA,2172
+nucliadb_utils-6.2.1.post3188.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+nucliadb_utils-6.2.1.post3188.dist-info/top_level.txt,sha256=fE3vJtALTfgh7bcAWcNhcfXkNPp_eVVpbKK-2IYua3E,15
+nucliadb_utils-6.2.1.post3188.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+nucliadb_utils-6.2.1.post3188.dist-info/RECORD,,