groundx 2.0.15__py3-none-any.whl → 2.7.7__py3-none-any.whl

groundx/extract/services/status.py

@@ -0,0 +1,197 @@
+import typing
+
+from fastapi import Response
+
+from .logger import Logger
+from ..settings.settings import ContainerSettings
+
+
+class Status:
+    def __init__(
+        self,
+        cfg: ContainerSettings,
+        logger: Logger,
+    ) -> None:
+        import redis
+
+        rl_port = 6379
+        rl_host = cfg.status_broker()
+        rl_ssl = False
+        if rl_host.endswith("/0"):
+            rl_host = rl_host[:-2]
+        if rl_host.startswith("redis://"):
+            rl_host = rl_host[8:]
+        elif rl_host.startswith("rediss://"):
+            rl_host = rl_host[9:]
+            rl_ssl = True
+        if ":" in rl_host:
+            base, number = rl_host.rsplit(":", 1)
+            if number.isdigit():
+                rl_port = int(number)
+                rl_host = base
+
+        self.client = redis.Redis(
+            host=rl_host, port=rl_port, decode_responses=True, ssl=rl_ssl
+        )
+        self.host = rl_host
+        self.port = rl_port
+
+        self.config = cfg
+        self.logger = logger
+
+        self.logger.info_msg(
+            f"\n\t[{self.config.service}] [status.Status.__init__]\n\t\t{self.host}:{self.port}",
+        )
+
+    def get_worker_state(
+        self, id: str, to: typing.Optional[int] = None
+    ) -> typing.Tuple[typing.Optional[int], int]:
+        online = self.client.get(self.key_worker_status(id))
+        if online is None or online == "offline":
+            return None, self.config.workers
+
+        key_worker_available = self.key_worker_available(id)
+
+        current_available = self.client.get(key_worker_available)
+        if current_available is None:
+            return None, self.config.workers
+
+        return int(current_available), self.config.workers  # type: ignore
+
+    def get_service_state(self) -> typing.Tuple[int, int]:
+        available = 0
+
+        keys: typing.Iterator[str] = self.client.scan_iter(  # type: ignore
+            match=f"{self.config.service}:*:requests",
+            count=1000,
+        )
+        for key in keys:
+            value = self.client.get(key)
+            if value is not None:
+                available += int(value)  # type: ignore
+
+        total = 0
+
+        keys: typing.Iterator[str] = self.client.scan_iter(  # type: ignore
+            match=f"{self.config.service}:*:total", count=1000
+        )
+        for key in keys:
+            value = self.client.get(key)
+            if value is not None:
+                total += int(value)  # type: ignore
+
+        if total == 0:
+            total = self.config.workers
+
+        return available, total
+
+    def key_worker_available(self, id: str) -> str:
+        return f"{self.config.service}:{id}:requests"
+
+    def key_worker_status(self, id: str) -> str:
+        return f"{self.config.service}:{id}:status"
+
+    def key_worker_total(self, id: str) -> str:
+        return f"{self.config.service}:{id}:total"
+
+    def refresh_worker(self, id: str, to: typing.Optional[int] = None) -> None:
+        self.refresh_worker_online(id, to)
+        self.refresh_worker_total(id, to)
+        self.refresh_worker_available(id, to)
+
+    def refresh_worker_available(
+        self, id: str, to: typing.Optional[int] = None
+    ) -> None:
+        key_worker_available = self.key_worker_available(id)
+        current_available = self.client.get(key_worker_available)
+        if current_available is None:
+            self.set_value(key_worker_available, self.config.workers, to)
+        else:
+            if to is not None:
+                if to > 0:
+                    self.client.expire(key_worker_available, to)
+            else:
+                self.client.expire(key_worker_available, self.config.cache_to)
+
+    def refresh_worker_online(self, id: str, to: typing.Optional[int] = None) -> None:
+        self.set_worker_online(id, to)
+
+    def refresh_worker_total(self, id: str, to: typing.Optional[int] = None) -> None:
+        self.set_value(self.key_worker_total(id), self.config.workers, to)
+
+    def set_headers(
+        self,
+        response: Response,
+        id: str,
+        available: typing.Optional[int],
+        total: typing.Optional[int],
+    ) -> typing.Any:
+        if available is None:
+            available = 0
+        if total is None:
+            total = 0
+
+        response.headers.update(
+            {
+                "X-RateLimit-Limit-Requests": str(total),
+                "X-RateLimit-Remaining-Requests": str(max(0, available)),
+                "X-ID": id,
+            }
+        )
+
+        return response
+
+    def set_value(
+        self, key: str, value: typing.Union[str, int], to: typing.Optional[int] = None
+    ) -> None:
+        if to is not None:
+            if to > 0:
+                self.client.set(key, value, ex=to)
+            else:
+                self.client.set(key, value, ex=self.config.cache_to)
+        else:
+            self.client.set(key, value, ex=self.config.cache_to)
+
+    def set_worker_available(self, id: str, to: typing.Optional[int] = None) -> None:
+        self.refresh_worker_online(id, to)
+
+        self.refresh_worker_total(id, to)
+
+        key_worker_available = self.key_worker_available(id)
+        current_available = self.client.get(key_worker_available)
+        if current_available is None:
+            current_available = self.config.workers
+            self.set_value(key_worker_available, current_available, to)
+        else:
+            self.set_value(
+                key_worker_available,
+                min(self.config.workers, int(current_available) + 1),  # type: ignore
+                to,
+            )
+
+    def set_worker_offline(self, id: str, to: typing.Optional[int] = None) -> None:
+        if to is None:
+            to = self.config.cache_to
+        self.logger.info_msg(f"\n\n\t\t[{self.config.service}] offline [{id}]\n")
+        self.set_value(self.key_worker_status(id), "offline", to)
+        self.set_worker_unavailable(id, to)
+
+    def set_worker_online(self, id: str, to: typing.Optional[int] = None) -> None:
+        self.set_value(self.key_worker_status(id), "online", to)
+
+    def set_worker_unavailable(self, id: str, to: typing.Optional[int] = None) -> None:
+        self.refresh_worker_online(id, to)
+
+        self.set_value(self.key_worker_total(id), self.config.workers, to)
+
+        key_worker_available = self.key_worker_available(id)
+        current_available = self.client.get(key_worker_available)
+        if current_available is None:
+            current_available = self.config.workers - 1
+            self.set_value(key_worker_available, current_available, to)
+        else:
+            self.set_value(
+                key_worker_available,
+                max(0, int(current_available) - 1),  # type: ignore
+                to,
+            )

groundx/extract/services/upload.py

@@ -0,0 +1,68 @@
+import typing
+
+from .logger import Logger
+from ..settings.settings import ContainerSettings
+
+
+@typing.runtime_checkable
+class UploadClient(typing.Protocol):
+    def get_object(self, url: str) -> typing.Optional[bytes]: ...
+
+    def put_object(
+        self,
+        bucket: str,
+        key: str,
+        data: bytes,
+        content_type: str = "application/octet-stream",
+    ) -> None: ...
+
+    def put_json_stream(
+        self,
+        bucket: str,
+        key: str,
+        data: bytes,
+        content_type: str = "application/octet-stream",
+    ) -> None: ...
+
+
+class Upload:
+    def __init__(
+        self,
+        settings: ContainerSettings,
+        logger: Logger,
+    ) -> None:
+        self.client: UploadClient
+        self.settings = settings
+        self.logger = logger
+
+        if self.settings.upload.type == "minio":
+            from .upload_minio import MinIOClient
+
+            self.client = MinIOClient(self.settings, self.logger)
+        elif self.settings.upload.type == "s3":
+            from .upload_s3 import S3Client
+
+            self.client = S3Client(self.settings, self.logger)
+        else:
+            raise Exception(f"unsupported upload.type [{self.settings.upload.type}]")
+
+    def get_object(self, url: str) -> typing.Optional[bytes]:
+        return self.client.get_object(url)
+
+    def put_object(
+        self,
+        bucket: str,
+        key: str,
+        data: bytes,
+        content_type: str = "application/octet-stream",
+    ) -> None:
+        self.client.put_object(bucket, key, data, content_type)
+
+    def put_json_stream(
+        self,
+        bucket: str,
+        key: str,
+        data: bytes,
+        content_type: str = "application/octet-stream",
+    ) -> None:
+        self.client.put_json_stream(bucket, key, data, content_type)

groundx/extract/services/upload_minio.py

@@ -0,0 +1,122 @@
+import typing
+
+from .logger import Logger
+from ..settings.settings import ContainerSettings
+
+
+class MinIOClient:
+    def __init__(
+        self,
+        settings: ContainerSettings,
+        logger: Logger,
+    ) -> None:
+        self.settings = settings
+        self.client = None
+        self.logger = logger
+        if self.settings.upload.type == "minio":
+            import json
+            from minio import Minio
+
+            self.client = Minio(
+                self.settings.upload.base_domain,
+                access_key=self.settings.upload.get_key(),
+                secret_key=self.settings.upload.get_secret(),
+                region=self.settings.upload.get_region(),
+                session_token=self.settings.upload.get_token(),
+                secure=self.settings.upload.ssl,
+            )
+
+            if not self.client.bucket_exists(self.settings.upload.bucket):
+                try:
+                    self.client.make_bucket(self.settings.upload.bucket)
+                    self.logger.info_msg(
+                        f"Bucket '{self.settings.upload.bucket}' created."
+                    )
+
+                    self.client.set_bucket_policy(
+                        self.settings.upload.bucket,
+                        json.dumps(
+                            {
+                                "Version": "2012-10-17",
+                                "Statement": [
+                                    {
+                                        "Effect": "Allow",
+                                        "Principal": {"AWS": ["*"]},
+                                        "Action": ["s3:GetObject"],
+                                        "Resource": [
+                                            f"arn:aws:s3:::{self.settings.upload.bucket}/*"
+                                        ],
+                                    }
+                                ],
+                            }
+                        ),
+                    )
+                except Exception as e:
+                    self.logger.warning_msg(str(e))
+                    self.logger.warning_msg(
+                        f"error creating bucket [{self.settings.upload.bucket}]"
+                    )
+
+    def get_object(self, url: str) -> typing.Optional[bytes]:
+        if not self.client:
+            return None
+
+        from minio.error import S3Error
+
+        try:
+            minio_uri_parts = url.replace("s3://", "").split("/")
+            bucket_name = minio_uri_parts[0]
+            object_name = "/".join(minio_uri_parts[1:])
+
+            response = self.client.get_object(bucket_name, object_name)
+
+            return response.read()
+        except S3Error as e:
+            self.logger.error_msg(f"Failed to get object from {url}: {str(e)}")
+            raise
+
+    def put_object(
+        self,
+        bucket: str,
+        key: str,
+        data: bytes,
+        content_type: str = "application/octet-stream",
+    ) -> None:
+        if not self.client:
+            return
+
+        import io
+
+        from minio.error import S3Error
+
+        try:
+            if isinstance(data, str):
+                data = data.encode("utf-8")
+
+            self.client.put_object(
+                bucket_name=bucket,
+                object_name=key,
+                data=io.BytesIO(data),
+                length=len(data),
+                content_type=content_type,
+            )
+        except S3Error as e:
+            self.logger.error_msg(f"Failed to put object in {bucket}/{key}: {str(e)}")
+            raise
+
+    def put_json_stream(
+        self,
+        bucket: str,
+        key: str,
+        data: bytes,
+        content_type: str = "application/octet-stream",
+    ) -> None:
+        if not self.client:
+            return
+
+        self.put_object(
+            bucket,
+            key,
+            data,
+            content_type,
+        )

groundx/extract/services/upload_s3.py

@@ -0,0 +1,91 @@
+import typing
+
+from .logger import Logger
+from ..settings.settings import ContainerSettings
+
+
+class S3Client:
+    def __init__(self, settings: ContainerSettings, logger: Logger) -> None:
+        self.settings = settings
+        self.client = None
+        self.logger = logger
+        if self.settings.upload.type == "s3":
+            import boto3, certifi
+            from botocore.config import Config
+
+            self.client = boto3.client(  # pyright: ignore[reportUnknownMemberType]
+                "s3",
+                aws_access_key_id=self.settings.upload.get_key(),
+                aws_secret_access_key=self.settings.upload.get_secret(),
+                aws_session_token=self.settings.upload.get_token(),
+                config=Config(max_pool_connections=50),
+                region_name=self.settings.upload.get_region(),
+                verify=certifi.where(),
+            )
+
+    def get_object(self, url: str) -> typing.Optional[bytes]:
+        if not self.client:
+            print("get_object no client")
+            return None
+
+        try:
+            if url.startswith("s3://"):
+                s3_uri_parts = url.replace("s3://", "").split("/")
+                s3_bucket = s3_uri_parts[0]
+                s3_key = "/".join(s3_uri_parts[1:])
+            else:
+                s3_bucket = self.settings.upload.bucket
+                s3_key = url
+                if url.startswith("/"):
+                    s3_key = url[1:]
+
+            response = self.client.get_object(Bucket=s3_bucket, Key=s3_key)
+
+            return response["Body"].read()
+        except Exception as e:
+            self.logger.error_msg(f"[{url}] exception: {e}")
+            raise
+
+    def put_object(
+        self,
+        bucket: str,
+        key: str,
+        data: bytes,
+        content_type: str = "application/octet-stream",
+    ) -> None:
+        if not self.client:
+            return
+
+        self.client.put_object(
+            Bucket=bucket,
+            Key=key,
+            Body=data,
+            ContentType=content_type,
+        )
+
+    def put_json_stream(
+        self,
+        bucket: str,
+        key: str,
+        data: bytes,
+        content_type: str = "application/octet-stream",
+    ) -> None:
+        if not self.client:
+            return
+
+        import io
+
+        json_stream = io.BytesIO()
+
+        if isinstance(data, str):
+            data = data.encode("utf-8")
+
+        json_stream.write(data)
+        json_stream.seek(0)
+
+        self.put_object(
+            bucket,
+            key,
+            json_stream.getvalue(),
+            content_type,
+        )

groundx/extract/services/utility.py

@@ -0,0 +1,52 @@
+import typing
+
+
+def get_config_path() -> typing.Optional[str]:
+    import sys
+
+    if "-c" in sys.argv:
+        config_index = sys.argv.index("-c") + 1
+        if config_index < len(sys.argv):
+            return sys.argv[config_index]
+
+    return None
+
+
+def get_gunicorn_threads() -> int:
+    import importlib.util
+
+    try:
+        conf_path = get_config_path()
+        if conf_path is not None:
+            spec = importlib.util.spec_from_file_location("gunicorn_conf", conf_path)
+            if spec and spec.loader:
+                gunicorn_conf = importlib.util.module_from_spec(spec)
+                spec.loader.exec_module(gunicorn_conf)
+                return gunicorn_conf.threads
+            return 0
+        return 1
+    except Exception:
+        return 1
+
+
+def get_thread_id(
+    thread_ids: typing.Dict[str, str],
+) -> typing.Tuple[str, typing.Dict[str, str]]:
+    import secrets, threading
+
+    thread_name = threading.current_thread().name
+    if thread_name not in thread_ids:
+        thread_ids[thread_name] = secrets.token_hex(4)
+    return thread_ids[thread_name], thread_ids
+
+
+def get_worker_id() -> str:
+    import os
+
+    from multiprocessing import current_process
+
+    name = os.environ.get("HOSTNAME")
+    if name is None or name == "":
+        return str(current_process().pid)
+
+    return name

groundx/extract/settings/__init__.py

@@ -0,0 +1,15 @@
+from .settings import (
+    AgentSettings,
+    ContainerSettings,
+    ContainerUploadSettings,
+    GroundXSettings,
+    GCP_CREDENTIALS,
+)
+
+__all__ = [
+    "AgentSettings",
+    "ContainerSettings",
+    "ContainerUploadSettings",
+    "GroundXSettings",
+    "GCP_CREDENTIALS",
+]