financechatbotkit 2.0.0__py3-none-any.whl

telebotkit/firestore/client.py

@@ -0,0 +1,141 @@
+"""Shared Firestore client bootstrap and error helpers."""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+_SERVICE_ACCOUNT_BASE_FIELDS = {
+    "type": "service_account",
+    "auth_uri": "https://accounts.google.com/o/oauth2/auth",
+    "token_uri": "https://oauth2.googleapis.com/token",
+    "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs",
+}
+
+
+class FirestoreQuotaError(Exception):
+    """Raised when a Firestore operation fails due to quota or write limits."""
+
+
+def raise_if_quota_error(error: Exception) -> None:
+    """Re-raise quota-related Google API exceptions in a stable app-level form."""
+
+    try:
+        from google.api_core import exceptions as google_api_exceptions  # type: ignore
+
+        if isinstance(error, google_api_exceptions.ResourceExhausted):
+            raise FirestoreQuotaError(str(error)) from error
+    except ImportError:
+        return
+
+
+def load_service_account_info_from_env() -> dict[str, Any] | None:
+    """Construct service-account credentials from discrete GOOGLE_* env vars."""
+
+    client_email = os.environ.get("GOOGLE_CLIENT_EMAIL", "").strip()
+    if not client_email:
+        return None
+
+    private_key = os.environ.get("GOOGLE_PRIVATE_KEY", "").strip()
+    if "\\n" in private_key and "\n" not in private_key:
+        private_key = private_key.replace("\\n", "\n")
+
+    return {
+        **_SERVICE_ACCOUNT_BASE_FIELDS,
+        "project_id": os.environ.get("GOOGLE_PROJECT_ID", "").strip(),
+        "private_key_id": os.environ.get("GOOGLE_PRIVATE_KEY_ID", "").strip(),
+        "private_key": private_key,
+        "client_email": client_email,
+        "client_id": os.environ.get("GOOGLE_CLIENT_ID", "").strip(),
+        "client_x509_cert_url": os.environ.get("GOOGLE_CLIENT_X509_CERT_URL", "").strip(),
+    }
+
+
+class FirestoreClientProvider:
+    """Lazily create and cache a Firestore client."""
+
+    def __init__(self) -> None:
+        self._client = None
+
+    def reset(self) -> None:
+        self._client = None
+
+    def get_client(self):
+        if self._client is None:
+            project_id = (
+                os.environ.get("GOOGLE_PROJECT_ID")
+                or os.environ.get("FIRESTORE_PROJECT_ID")
+                or None
+            )
+
+            from google.cloud import firestore  # type: ignore
+
+            service_account_info = load_service_account_info_from_env()
+            if service_account_info:
+                from google.oauth2 import service_account  # type: ignore
+
+                credentials = service_account.Credentials.from_service_account_info(
+                    service_account_info,
+                    scopes=["https://www.googleapis.com/auth/datastore"],
+                )
+                self._client = firestore.Client(
+                    project=project_id or service_account_info.get("project_id") or None,
+                    credentials=credentials,
+                )
+                logger.info("Firestore client created from individual GOOGLE_* env vars.")
+            else:
+                credentials_json = os.environ.get(
+                    "GOOGLE_APPLICATION_CREDENTIALS_JSON",
+                    "",
+                ).strip()
+                if credentials_json:
+                    from google.oauth2 import service_account  # type: ignore
+
+                    info = json.loads(credentials_json)
+                    credentials = service_account.Credentials.from_service_account_info(
+                        info,
+                        scopes=["https://www.googleapis.com/auth/datastore"],
+                    )
+                    self._client = firestore.Client(
+                        project=project_id,
+                        credentials=credentials,
+                    )
+                    logger.info(
+                        "Firestore client created from GOOGLE_APPLICATION_CREDENTIALS_JSON."
+                    )
+                else:
+                    self._client = firestore.Client(project=project_id)
+                    logger.info(
+                        "Firestore client created from Application Default Credentials."
+                    )
+        return self._client
+
+FirestoreClientFactory = FirestoreClientProvider
+
+_DEFAULT_CLIENT_PROVIDER = FirestoreClientProvider()
+
+
+def get_client():
+    """Return the shared cached Firestore client."""
+
+    return _DEFAULT_CLIENT_PROVIDER.get_client()
+
+
+def reset_client() -> None:
+    """Reset the shared cached Firestore client."""
+
+    _DEFAULT_CLIENT_PROVIDER.reset()
+
+__all__ = [
+    "FirestoreClientFactory",
+    "FirestoreClientProvider",
+    "FirestoreQuotaError",
+    "get_client",
+    "load_service_account_info_from_env",
+    "raise_if_quota_error",
+    "reset_client",
+]

telebotkit/firestore/documents.py

@@ -0,0 +1,164 @@
+"""Typed Firestore repository helpers."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any, Generic, TypeVar
+
+from telebotkit.firestore.client import get_client
+from telebotkit.firestore.fetch import FetchResult, fetch_document, invalidate_document_cache
+
+DocumentT = TypeVar("DocumentT")
+
+
+class DocumentStore(Generic[DocumentT]):
+    """Typed access to a Firestore collection with explicit document IDs."""
+
+    def __init__(
+        self,
+        *,
+        collection_name: str,
+        from_dict: Callable[[dict[str, Any]], DocumentT],
+        to_dict: Callable[[DocumentT], dict[str, Any]],
+        client_provider: Callable[[], Any] = get_client,
+        cache_ttl_seconds: float | None = None,
+    ) -> None:
+        self._collection_name = collection_name
+        self._from_dict = from_dict
+        self._to_dict = to_dict
+        self._client_provider = client_provider
+        self._cache_ttl_seconds = cache_ttl_seconds
+
+    def get_result(
+        self,
+        document_id: str,
+        *,
+        force_refresh: bool = False,
+    ) -> FetchResult[DocumentT]:
+        return fetch_document(
+            self._collection_name,
+            document_id,
+            parse=self._from_dict,
+            client_provider=self._client_provider,
+            cache_ttl_seconds=self._cache_ttl_seconds,
+            force_refresh=force_refresh,
+        )
+
+    def get(
+        self,
+        document_id: str,
+        *,
+        force_refresh: bool = False,
+    ) -> DocumentT | None:
+        result = self.get_result(document_id, force_refresh=force_refresh)
+        if result.error is not None:
+            raise result.error
+        return result.data
+
+    def get_or_create(
+        self,
+        document_id: str,
+        factory: Callable[[], DocumentT],
+    ) -> DocumentT:
+        document = self.get(document_id)
+        return document if document is not None else factory()
+
+    def save(
+        self,
+        document_id: str,
+        value: DocumentT,
+        *,
+        merge: bool = False,
+    ) -> None:
+        self._collection().document(document_id).set(self._to_dict(value), merge=merge)
+        invalidate_document_cache(self._collection_name, document_id)
+
+    def update(
+        self,
+        document_id: str,
+        fields: dict[str, Any],
+        *,
+        merge: bool = True,
+    ) -> None:
+        self._collection().document(document_id).set(fields, merge=merge)
+        invalidate_document_cache(self._collection_name, document_id)
+
+    def set_fields(
+        self,
+        document_id: str,
+        payload: dict[str, Any],
+        *,
+        merge: bool = True,
+    ) -> None:
+        self.update(document_id, payload, merge=merge)
+
+    def list_all(self) -> list[DocumentT]:
+        return [
+            self._from_dict(document.to_dict() or {})
+            for document in self._collection().stream()
+        ]
+
+    def document(self, document_id: str):
+        return self._collection().document(document_id)
+
+    def _collection(self):
+        return self._client_provider().collection(self._collection_name)
+
+
+class SharedDocumentStore(Generic[DocumentT]):
+    """Typed repository for a single shared Firestore document."""
+
+    def __init__(
+        self,
+        *,
+        collection_name: str,
+        document_id: str,
+        from_dict: Callable[[dict[str, Any]], DocumentT],
+        to_dict: Callable[[DocumentT], dict[str, Any]],
+        client_provider: Callable[[], Any] = get_client,
+        cache_ttl_seconds: float | None = None,
+    ) -> None:
+        self._store = DocumentStore(
+            collection_name=collection_name,
+            from_dict=from_dict,
+            to_dict=to_dict,
+            client_provider=client_provider,
+            cache_ttl_seconds=cache_ttl_seconds,
+        )
+        self._document_id = document_id
+
+    def get_result(
+        self,
+        *,
+        force_refresh: bool = False,
+    ) -> FetchResult[DocumentT]:
+        return self._store.get_result(
+            self._document_id,
+            force_refresh=force_refresh,
+        )
+
+    def get(self, *, force_refresh: bool = False) -> DocumentT | None:
+        return self._store.get(
+            self._document_id,
+            force_refresh=force_refresh,
+        )
+
+    def get_or_create(self, factory: Callable[[], DocumentT]) -> DocumentT:
+        return self._store.get_or_create(self._document_id, factory)
+
+    def save(self, value: DocumentT, *, merge: bool = False) -> None:
+        self._store.save(self._document_id, value, merge=merge)
+
+    def update(self, fields: dict[str, Any], *, merge: bool = True) -> None:
+        self._store.update(self._document_id, fields, merge=merge)
+
+    def set_fields(self, payload: dict[str, Any], *, merge: bool = True) -> None:
+        self.update(payload, merge=merge)
+
+    def document(self):
+        return self._store.document(self._document_id)
+
+__all__ = [
+    "DocumentStore",
+    "SharedDocumentStore",
+]

telebotkit/firestore/fetch.py

@@ -0,0 +1,228 @@
+"""Higher-level Firestore document fetch helpers with optional caching."""
+
+from __future__ import annotations
+
+from copy import deepcopy
+from dataclasses import dataclass
+import time
+from typing import Any, Callable, Generic, TypeVar
+
+from telebotkit.firestore.client import FirestoreQuotaError, get_client, raise_if_quota_error
+
+DocumentT = TypeVar("DocumentT")
+
+
+@dataclass(frozen=True)
+class FetchResult(Generic[DocumentT]):
+    """Result of fetching one Firestore document."""
+
+    collection_name: str
+    document_id: str
+    found: bool
+    data: DocumentT | None = None
+    raw_data: dict[str, Any] | None = None
+    error: Exception | None = None
+    from_cache: bool = False
+
+    @property
+    def ok(self) -> bool:
+        return self.error is None
+
+    @property
+    def missing(self) -> bool:
+        return self.error is None and not self.found
+
+    @property
+    def has_data(self) -> bool:
+        return self.error is None and self.found and self.data is not None
+
+    @property
+    def is_quota_error(self) -> bool:
+        return isinstance(self.error, FirestoreQuotaError)
+
+
+@dataclass
+class _CacheEntry:
+    found: bool
+    raw_data: dict[str, Any] | None
+    expires_at: float
+
+
+_DOCUMENT_CACHE: dict[tuple[str, str], _CacheEntry] = {}
+
+
+def clear_document_cache() -> None:
+    """Drop all cached document fetch entries."""
+
+    _DOCUMENT_CACHE.clear()
+
+
+def invalidate_document_cache(collection_name: str, document_id: str) -> None:
+    """Drop one cached document fetch entry."""
+
+    _DOCUMENT_CACHE.pop((collection_name, document_id), None)
+
+
+def fetch_document(
+    collection_name: str,
+    document_id: str,
+    *,
+    parse: Callable[[dict[str, Any]], DocumentT] | None = None,
+    client_provider: Callable[[], Any] = get_client,
+    cache_ttl_seconds: float | None = None,
+    force_refresh: bool = False,
+) -> FetchResult[Any]:
+    """Fetch one Firestore document as raw data or a parsed value."""
+
+    cache_key = (collection_name, document_id)
+    if not force_refresh and cache_ttl_seconds and cache_ttl_seconds > 0:
+        cached_entry = _DOCUMENT_CACHE.get(cache_key)
+        if cached_entry is not None:
+            if cached_entry.expires_at > time.monotonic():
+                return _result_from_entry(
+                    collection_name=collection_name,
+                    document_id=document_id,
+                    entry=cached_entry,
+                    parse=parse,
+                    from_cache=True,
+                )
+            invalidate_document_cache(collection_name, document_id)
+
+    try:
+        document = client_provider().collection(collection_name).document(document_id).get()
+        if not document.exists:
+            _cache_entry(
+                cache_key,
+                found=False,
+                raw_data=None,
+                cache_ttl_seconds=cache_ttl_seconds,
+            )
+            return FetchResult(
+                collection_name=collection_name,
+                document_id=document_id,
+                found=False,
+            )
+
+        raw_data = dict(document.to_dict() or {})
+        _cache_entry(
+            cache_key,
+            found=True,
+            raw_data=raw_data,
+            cache_ttl_seconds=cache_ttl_seconds,
+        )
+        return _build_result(
+            collection_name=collection_name,
+            document_id=document_id,
+            found=True,
+            raw_data=raw_data,
+            parse=parse,
+            from_cache=False,
+        )
+    except Exception as error:
+        normalized_error = _normalize_error(error)
+        return FetchResult(
+            collection_name=collection_name,
+            document_id=document_id,
+            found=False,
+            error=normalized_error,
+        )
+
+
+def _result_from_entry(
+    *,
+    collection_name: str,
+    document_id: str,
+    entry: _CacheEntry,
+    parse: Callable[[dict[str, Any]], DocumentT] | None,
+    from_cache: bool,
+) -> FetchResult[Any]:
+    return _build_result(
+        collection_name=collection_name,
+        document_id=document_id,
+        found=entry.found,
+        raw_data=deepcopy(entry.raw_data),
+        parse=parse,
+        from_cache=from_cache,
+    )
+
+
+def _build_result(
+    *,
+    collection_name: str,
+    document_id: str,
+    found: bool,
+    raw_data: dict[str, Any] | None,
+    parse: Callable[[dict[str, Any]], DocumentT] | None,
+    from_cache: bool,
+) -> FetchResult[Any]:
+    if not found:
+        return FetchResult(
+            collection_name=collection_name,
+            document_id=document_id,
+            found=False,
+            from_cache=from_cache,
+        )
+
+    safe_raw = deepcopy(raw_data) if raw_data is not None else {}
+    if parse is None:
+        return FetchResult(
+            collection_name=collection_name,
+            document_id=document_id,
+            found=True,
+            data=safe_raw,
+            raw_data=safe_raw,
+            from_cache=from_cache,
+        )
+
+    try:
+        parsed = parse(deepcopy(safe_raw))
+    except Exception as error:
+        return FetchResult(
+            collection_name=collection_name,
+            document_id=document_id,
+            found=True,
+            raw_data=safe_raw,
+            error=error,
+            from_cache=from_cache,
+        )
+
+    return FetchResult(
+        collection_name=collection_name,
+        document_id=document_id,
+        found=True,
+        data=parsed,
+        raw_data=safe_raw,
+        from_cache=from_cache,
+    )
+
+
+def _cache_entry(
+    cache_key: tuple[str, str],
+    *,
+    found: bool,
+    raw_data: dict[str, Any] | None,
+    cache_ttl_seconds: float | None,
+) -> None:
+    if not cache_ttl_seconds or cache_ttl_seconds <= 0:
+        return
+    _DOCUMENT_CACHE[cache_key] = _CacheEntry(
+        found=found,
+        raw_data=deepcopy(raw_data),
+        expires_at=time.monotonic() + cache_ttl_seconds,
+    )
+
+
+def _normalize_error(error: Exception) -> Exception:
+    try:
+        raise_if_quota_error(error)
+    except Exception as normalized_error:
+        return normalized_error
+    return error
+
+
+__all__ = [
+    "FetchResult",
+    "clear_document_cache",
+    "fetch_document",
+    "invalidate_document_cache",
+]

telebotkit/firestore/locks.py

@@ -0,0 +1,74 @@
+"""Firestore-backed short-lived lease helpers."""
+
+from __future__ import annotations
+
+import logging
+import time
+from typing import Any
+
+from telebotkit.firestore.client import get_client
+
+logger = logging.getLogger(__name__)
+
+
+class LeaseStore:
+    """Acquire and renew short-lived Firestore document leases."""
+
+    def __init__(
+        self,
+        *,
+        collection_name: str,
+        client_provider=get_client,
+    ) -> None:
+        self._collection_name = collection_name
+        self._client_provider = client_provider
+
+    def acquire(
+        self,
+        *,
+        lease_name: str,
+        owner_id: str,
+        ttl_seconds: float,
+    ) -> bool:
+        from google.cloud import firestore  # type: ignore
+
+        client = self._client_provider()
+        document_reference = client.collection(self._collection_name).document(lease_name)
+        transaction = client.transaction()
+        now = time.time()
+        expires_at = now + max(ttl_seconds, 1.0)
+
+        @firestore.transactional
+        def _acquire(transaction_context: Any) -> bool:
+            snapshot = document_reference.get(transaction=transaction_context)
+            if snapshot.exists:
+                data = snapshot.to_dict() or {}
+                current_owner = str(data.get("owner_id") or "")
+                current_expires_at = float(data.get("expires_at") or 0.0)
+                if current_owner and current_owner != owner_id and current_expires_at > now:
+                    return False
+            transaction_context.set(
+                document_reference,
+                {
+                    "owner_id": owner_id,
+                    "expires_at": expires_at,
+                    "updated_at": now,
+                },
+            )
+            return True
+
+        return bool(_acquire(transaction))
+
+    def try_acquire(
+        self,
+        *,
+        lease_name: str,
+        owner_id: str,
+        ttl_seconds: float,
+    ) -> bool:
+        return self.acquire(
+            lease_name=lease_name,
+            owner_id=owner_id,
+            ttl_seconds=ttl_seconds,
+        )
+__all__ = ["LeaseStore"]

telebotkit/firestore/upload.py

@@ -0,0 +1,75 @@
+"""Firestore upload helpers for prevalidated JSON payloads."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from telebotkit.firestore.client import get_client
+from telebotkit.sheets import load_typed_rows_json, validate_typed_rows_payload
+
+
+def upload_typed_rows_payload(
+    *,
+    payload: dict[str, Any],
+    collection_name: str,
+    document_id: str,
+    document_type: str = "typed_rows",
+    records_field: str = "records",
+    metadata_field: str = "metadata",
+    key_field_name: str = "key",
+) -> str:
+    validated = validate_typed_rows_payload(
+        payload,
+        document_type=document_type,
+        records_field=records_field,
+        metadata_field=metadata_field,
+        key_field_name=key_field_name,
+    )
+    get_client().collection(collection_name).document(document_id).set(
+        {
+            "document_type": document_type,
+            "schema_version": int(validated.get("schema_version", 1) or 1),
+            records_field: validated.get(records_field) or {},
+            metadata_field: validated.get(metadata_field) or {},
+        },
+        merge=True,
+    )
+    meta = validated.get(metadata_field) or {}
+    return (
+        "업로드 완료:\n"
+        f"- Firestore: {collection_name}/{document_id}\n"
+        f"- record_count: {meta.get('record_count')}\n"
+        f"- size_bytes: {meta.get('size_bytes')}\n"
+    )
+
+
+def upload_typed_rows_json(
+    *,
+    json_path: str,
+    collection_name: str,
+    document_id: str,
+    document_type: str = "typed_rows",
+    records_field: str = "records",
+    metadata_field: str = "metadata",
+    key_field_name: str = "key",
+) -> str:
+    payload = load_typed_rows_json(
+        json_path,
+        document_type=document_type,
+        records_field=records_field,
+        metadata_field=metadata_field,
+        key_field_name=key_field_name,
+    )
+    result = upload_typed_rows_payload(
+        payload=payload,
+        collection_name=collection_name,
+        document_id=document_id,
+        document_type=document_type,
+        records_field=records_field,
+        metadata_field=metadata_field,
+        key_field_name=key_field_name,
+    )
+    return result.replace("업로드 완료:\n", f"업로드 완료:\n- JSON: {json_path}\n", 1)
+
+
+__all__ = ["upload_typed_rows_json", "upload_typed_rows_payload"]