groundx 2.4.5__py3-none-any.whl → 2.4.10__py3-none-any.whl

groundx/extract/services/sheets_client.py

@@ -0,0 +1,160 @@
+import json, os, typing
+from pathlib import Path
+
+from google.oauth2 import service_account
+from googleapiclient.discovery import (
+    build,  # pyright: ignore[reportUnknownVariableType]
+)
+import gspread
+
+from ..classes.settings import ContainerSettings, GCP_CREDENTIALS
+
+SPREADSHEET_MIME = "application/vnd.google-apps.spreadsheet"
+
+
+class SheetsClient:
+    client: gspread.Client
+    # drive: DriveResource
+    drive: typing.Any
+    settings: ContainerSettings
+
+    def __init__(
+        self,
+        settings: ContainerSettings,
+        scopes: typing.Optional[typing.List[str]] = None,
+    ):
+        self.scopes: typing.List[str] = scopes or [
+            "https://www.googleapis.com/auth/spreadsheets",
+            "https://www.googleapis.com/auth/drive",
+        ]
+        self.settings = settings
+
+        creds_dict = _load_credentials_from_env()
+        if not creds_dict:
+            raise ValueError(f"{GCP_CREDENTIALS} does not load valid credentials")
+
+        creds = service_account.Credentials.from_service_account_info(
+            creds_dict, scopes=scopes
+        )
+        self.drive = build("drive", "v3", credentials=creds)
+
+        auth_scopes = self.scopes
+        if scopes:
+            auth_scopes = scopes
+
+        self.client = gspread.service_account_from_dict(creds_dict, scopes=auth_scopes)
+
+    def create_headers_if_missing(
+        self, ws: gspread.Worksheet, headers: typing.List[str]
+    ) -> None:
+        existing = ws.row_values(1)
+        if not existing:
+            ws.insert_row(headers, 1)
+
+    def find_sheet_by_name(
+        self,
+        spreadsheet_name: str,
+        drive_id: str,
+    ) -> typing.Optional[str]:
+        cln = spreadsheet_name.replace("'", "\\'")
+
+        q = f"name = '{cln}' and mimeType = '{SPREADSHEET_MIME}' and trashed = false"
+
+        resp = (
+            self.drive.files()
+            .list(
+                q=q,
+                corpora="drive",
+                driveId=drive_id,
+                includeItemsFromAllDrives=True,
+                supportsAllDrives=True,
+                fields="files(id, name)",
+            )
+            .execute()
+        )
+        files = resp.get("files", [])
+        return files[0].get("id") if files else None
+
+    def open_or_create_spreadsheet(
+        self,
+        spreadsheet_name: str,
+        drive_id: str,
+        sheet_1_title: typing.Optional[str] = None,
+    ) -> gspread.Spreadsheet:
+        file_id = self.find_sheet_by_name(spreadsheet_name, drive_id)
+        if file_id:
+            return self.client.open_by_key(file_id)
+
+        if self.settings.google_sheets_template_id:
+            created = (
+                self.drive.files()
+                .copy(
+                    fileId=self.settings.google_sheets_template_id,
+                    body={"name": spreadsheet_name, "parents": [drive_id]},
+                    supportsAllDrives=True,
+                    fields="id,name,parents,driveId",
+                )
+                .execute()
+            )
+        else:
+            created = (
+                self.drive.files()
+                .create(
+                    body={
+                        "name": spreadsheet_name,
+                        "mimeType": SPREADSHEET_MIME,
+                        "parents": [drive_id],
+                    },
+                    supportsAllDrives=True,
+                    fields="id,name,parents,driveId",
+                )
+                .execute()
+            )
+
+        cid = created.get("id")
+        if not cid:
+            raise Exception(f"create spreadsheet failed\n{created}")
+
+        sh = self.client.open_by_key(cid)
+
+        if sheet_1_title:
+            sh.sheet1.update_title(sheet_1_title)
+
+        return sh
+
+    def open_or_create_worksheet(
+        self,
+        sh: gspread.Spreadsheet,
+        title: str,
+        headers: typing.List[str],
+        rows: int = 1000,
+    ) -> gspread.Worksheet:
+        cols = len(headers)
+        try:
+            ws = sh.worksheet(title)
+            self.create_headers_if_missing(ws, headers)
+        except gspread.WorksheetNotFound:
+            ws = sh.add_worksheet(title=title, rows=rows, cols=cols)
+            ws.append_row(headers)
+
+        return ws
+
+
+def _load_credentials_from_env() -> typing.Optional[typing.Dict[str, typing.Any]]:
+    raw = os.environ.get(GCP_CREDENTIALS)
+    if not raw:
+        if Path("./gcv.json").exists():
+            with open("./gcv.json") as f:
+                data = f.read()
+            return json.loads(data)
+
+        return None
+
+    try:
+        creds = json.loads(raw)
+        if not isinstance(creds, dict):
+            raise ValueError(f"{GCP_CREDENTIALS} is not type dict [{type(creds)}]")
+
+        return typing.cast(typing.Dict[str, typing.Any], creds)
+    except Exception as e:
+        raise ValueError(f"{GCP_CREDENTIALS} is set but not valid JSON: {e}") from e

groundx/extract/services/status.py

@@ -0,0 +1,197 @@
+import typing
+
+from fastapi import Response
+
+from ..classes.settings import ContainerSettings
+from .logger import Logger
+
+
+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,73 @@
+import typing
+
+from ..classes.settings import ContainerSettings
+from .logger import Logger
+
+
+@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_file(self, url: str) -> bytes:
+        return bytes()
+
+    def get_object(self, url: str) -> typing.Optional[bytes]:
+        self.client.get_object(url)
+
+        return None
+
+    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 ..classes.settings import ContainerSettings
+from .logger import Logger
+
+
+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,84 @@
+import typing
+
+from ..classes.settings import ContainerSettings
+from .logger import Logger
+
+
+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:
+            return None
+
+        try:
+            s3_uri_parts = url.replace("s3://", "").split("/")
+            s3_bucket = s3_uri_parts[0]
+            s3_key = "/".join(s3_uri_parts[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