altcodepro-polydb-python 2.2.2__py3-none-any.whl → 2.2.4__py3-none-any.whl

polydb/adapters/GCPPubSubAdapter.py

@@ -0,0 +1,217 @@
+import os
+import json
+import threading
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+from google.cloud import pubsub_v1
+from google.api_core.exceptions import AlreadyExists, NotFound
+
+from ..base.QueueAdapter import QueueAdapter
+from ..errors import ConnectionError, QueueError
+from ..retry import retry
+from ..json_safe import json_safe
+
+
+JsonLike = Union[Dict[str, Any], List[Any], str, int, float, bool, None]
+
+
+class GCPPubSubAdapter(QueueAdapter):
+    """
+    Production-grade GCP Pub/Sub adapter.
+
+    - Emulator support via PUBSUB_EMULATOR_HOST (google client honors it)
+    - Auto topic + subscription creation
+    - send accepts Any (tests send str)
+    - receive returns [{"id","ack_id","body"}...]
+    - ack() method for tests
+    """
+
+    def __init__(
+        self,
+        project_id: Optional[str] = None,
+        topic: Optional[str] = None,
+        subscription: Optional[str] = None,
+    ):
+        super().__init__()
+
+        self.project_id: str = project_id or os.getenv("GOOGLE_CLOUD_PROJECT", "polydb-test")
+        self.default_topic: str = topic or os.getenv("PUBSUB_TOPIC", "polydb-topic")
+        self.default_subscription: str = subscription or os.getenv(
+            "PUBSUB_SUBSCRIPTION", "polydb-sub"
+        )
+
+        self._publisher: Optional[pubsub_v1.PublisherClient] = None
+        self._subscriber: Optional[pubsub_v1.SubscriberClient] = None
+        self._lock = threading.Lock()
+
+        self._initialize_clients()
+
+    def _initialize_clients(self) -> None:
+        try:
+            with self._lock:
+                if self._publisher and self._subscriber:
+                    return
+                # google client auto-detects PUBSUB_EMULATOR_HOST if set
+                self._publisher = pubsub_v1.PublisherClient()
+                self._subscriber = pubsub_v1.SubscriberClient()
+                self.logger.info("Initialized Pub/Sub clients")
+        except Exception as e:
+            raise ConnectionError(f"Failed to initialize Pub/Sub: {e}")
+
+    def _resolve_names(self, queue_name: str) -> Tuple[str, str]:
+        """
+        Your tests call send() / receive() without queue_name,
+        so they pass default="default". That should map to adapter defaults.
+        """
+        if not queue_name or queue_name == "default":
+            return self.default_topic, self.default_subscription
+        # If user passes a custom name, use it for both (simple convention)
+        return queue_name, queue_name
+
+    def _topic_path(self, topic: str) -> str:
+        if not self._publisher:
+            raise ConnectionError("Pub/Sub publisher not initialized")
+        return self._publisher.topic_path(self.project_id, topic)
+
+    def _subscription_path(self, subscription: str) -> str:
+        if not self._subscriber:
+            raise ConnectionError("Pub/Sub subscriber not initialized")
+        return self._subscriber.subscription_path(self.project_id, subscription)
+
+    def _ensure_topic(self, topic: str) -> str:
+        if not self._publisher:
+            raise ConnectionError("Pub/Sub publisher not initialized")
+        path = self._topic_path(topic)
+        try:
+            self._publisher.get_topic(request={"topic": path})
+        except NotFound:
+            try:
+                self._publisher.create_topic(request={"name": path})
+                self.logger.info(f"Created Pub/Sub topic: {topic}")
+            except AlreadyExists:
+                pass
+        return path
+
+    def _ensure_subscription(self, topic: str, subscription: str) -> str:
+        if not self._subscriber:
+            raise ConnectionError("Pub/Sub subscriber not initialized")
+        if not self._publisher:
+            raise ConnectionError("Pub/Sub publisher not initialized")
+
+        topic_path = self._ensure_topic(topic)
+        sub_path = self._subscription_path(subscription)
+
+        try:
+            self._subscriber.get_subscription(request={"subscription": sub_path})
+        except NotFound:
+            try:
+                self._subscriber.create_subscription(
+                    request={"name": sub_path, "topic": topic_path}
+                )
+                self.logger.info(f"Created Pub/Sub subscription: {subscription}")
+            except AlreadyExists:
+                pass
+
+        return sub_path
+
+    # -------------------------
+    # API
+    # -------------------------
+
+    @retry(max_attempts=3, delay=1.0, exceptions=(QueueError,))
+    def send(self, message: JsonLike, queue_name: str = "default") -> str:
+        """
+        Publish message to Pub/Sub.
+        Tests send a string; production may send dict.
+        We envelope it as {"body": <message>}.
+        """
+        try:
+            if not self._publisher:
+                raise ConnectionError("Pub/Sub publisher not initialized")
+
+            topic, subscription = self._resolve_names(queue_name)
+            topic_path = self._ensure_topic(topic)
+            # Ensure subscription exists so receive works immediately in emulator
+            self._ensure_subscription(topic, subscription)
+
+            payload = {"body": message}
+            data = json.dumps(payload, default=json_safe).encode("utf-8")
+
+            future = self._publisher.publish(topic_path, data=data)
+            msg_id = future.result(timeout=10)
+
+            self.logger.debug(f"Published Pub/Sub message {msg_id} (topic={topic})")
+            return msg_id
+
+        except Exception as e:
+            raise QueueError(f"Pub/Sub send failed: {e}")
+
+    @retry(max_attempts=3, delay=1.0, exceptions=(QueueError,))
+    def receive(self, queue_name: str = "default", max_messages: int = 1) -> List[Dict[str, Any]]:
+        """
+        Pull messages.
+        IMPORTANT: Do NOT auto-ack here (so test_ack_message can ack explicitly).
+        """
+        try:
+            if not self._subscriber:
+                raise ConnectionError("Pub/Sub subscriber not initialized")
+
+            topic, subscription = self._resolve_names(queue_name)
+            sub_path = self._ensure_subscription(topic, subscription)
+
+            resp = self._subscriber.pull(
+                request={"subscription": sub_path, "max_messages": max_messages},
+                timeout=5,
+            )
+
+            out: List[Dict[str, Any]] = []
+            for r in resp.received_messages:
+                body: Any = None
+                try:
+                    decoded = json.loads(r.message.data.decode("utf-8"))
+                    if isinstance(decoded, dict) and "body" in decoded:
+                        body = decoded["body"]
+                    else:
+                        body = decoded
+                except Exception:
+                    body = r.message.data.decode("utf-8", errors="replace")
+
+                out.append(
+                    {
+                        "id": r.message.message_id,
+                        "ack_id": r.ack_id,
+                        "body": body,
+                    }
+                )
+
+            return out
+
+        except Exception as e:
+            raise QueueError(f"Pub/Sub receive failed: {e}")
+
+    def delete(self, message_id: str, queue_name: str = "default", pop_receipt: str = "") -> bool:
+        """
+        Pub/Sub delete == ack.
+        - pop_receipt maps to ack_id (preferred).
+        - message_id fallback kept for backward compatibility.
+        """
+        try:
+            if not self._subscriber:
+                raise ConnectionError("Pub/Sub subscriber not initialized")
+
+            _, subscription = self._resolve_names(queue_name)
+            sub_path = self._subscription_path(subscription)
+
+            ack_id = pop_receipt or message_id
+            if not ack_id:
+                raise QueueError("Pub/Sub delete requires ack_id (use pop_receipt from receive())")
+
+            self._subscriber.acknowledge(request={"subscription": sub_path, "ack_ids": [ack_id]})
+            return True
+
+        except Exception as e:
+            raise QueueError(f"Pub/Sub delete failed: {e}")
+
+    # Alias that your tests want
+    def ack(self, ack_id: str, queue_name: str = "default") -> bool:
+        return self.delete(message_id="", queue_name=queue_name, pop_receipt=ack_id)

polydb/adapters/GCPStorageAdapter.py

@@ -1,81 +1,226 @@
 # src/polydb/adapters/GCPStorageAdapter.py
+
+import mimetypes
 import os
 import threading
-from typing import List, Optional, cast
-from google.cloud.firestore import DocumentSnapshot
+from typing import Any, Dict, List, Optional
+
+from google.cloud import storage
+from google.api_core.exceptions import NotFound
+
 from ..base.ObjectStorageAdapter import ObjectStorageAdapter
 from ..errors import StorageError, ConnectionError
 from ..retry import retry
 
 
 class GCPStorageAdapter(ObjectStorageAdapter):
-    """GCP Cloud Storage with client reuse"""
+    """
+    Production-grade Google Cloud Storage adapter.
+
+    Features
+    --------
+    - Thread-safe client initialization
+    - Automatic bucket creation
+    - Emulator support (fake-gcs-server)
+    - Retry support
+    - Structured logging
+    """
 
-    def __init__(self):
+    def __init__(self, bucket_name: Optional[str] = None):
         super().__init__()
-        self.bucket_name = os.getenv("GCS_BUCKET_NAME", "default")
-        self._client = None
-        self._bucket = None
+
+        self.bucket_name: str = bucket_name or os.getenv("GCS_BUCKET_NAME", "default")
+        self.project_id: str = os.getenv("GOOGLE_CLOUD_PROJECT", "polydb-test")
+
+        # Emulator support
+        self._endpoint: Optional[str] = os.getenv("GCS_ENDPOINT")
+
+        self._client: Optional[storage.Client] = None
+        self._bucket: Optional[storage.Bucket] = None
+
         self._lock = threading.Lock()
+
         self._initialize_client()
 
-    def _initialize_client(self):
-        """Initialize GCS client once"""
-        try:
-            from google.cloud import storage
+    # ------------------------------------------------------------------
+    # Client initialization
+    # ------------------------------------------------------------------
 
+    def _initialize_client(self) -> None:
+        """Initialize GCS client once (thread-safe)"""
+        try:
             with self._lock:
-                if not self._client:
-                    self._client = storage.Client()
-                    self._bucket = self._client.bucket(self.bucket_name)
-                    self.logger.info("Initialized GCS client")
+                if self._client:
+                    return
+
+                if self._endpoint:
+                    self.logger.info(f"Using GCS emulator: {self._endpoint}")
+                    self._client = storage.Client(
+                        project=self.project_id,
+                        client_options={"api_endpoint": self._endpoint},
+                    )
+                else:
+                    self._client = storage.Client(project=self.project_id)
+
+                self._bucket = self._client.bucket(self.bucket_name)
+
+                # Ensure bucket exists
+                try:
+                    if not self._bucket.exists():
+                        self._bucket = self._client.create_bucket(self.bucket_name)
+                        self.logger.info(f"Created GCS bucket: {self.bucket_name}")
+                except Exception:
+                    # fake-gcs-server does not support bucket.exists()
+                    pass
+
+                self.logger.info(
+                    f"GCS initialized (bucket={self.bucket_name}, project={self.project_id})"
+                )
+
         except Exception as e:
             raise ConnectionError(f"Failed to initialize GCS: {str(e)}")
 
+    # ------------------------------------------------------------------
+    # Put object
+    # ------------------------------------------------------------------
+
     @retry(max_attempts=3, delay=1.0, exceptions=(StorageError,))
-    def _put_raw(self, key: str, data: bytes) -> str:
-        """Store object"""
+    def _put_raw(
+        self,
+        key: str,
+        data: bytes,
+        fileName: str = "",
+        media_type: Optional[str] = None,
+        metadata: Dict[str, Any] | None = None,
+    ) -> str:
+        """Upload object to GCS with media type, metadata, filename handling"""
         try:
-            if self._bucket:
-                blob = self._bucket.blob(key)
-                blob.upload_from_string(data)
-                self.logger.debug(f"Uploaded blob: {key}")
-            return key
+            if not self._bucket:
+                raise ConnectionError("GCS bucket not initialized")
+
+            # --------------------------------------------------
+            # Resolve filename
+            # --------------------------------------------------
+            filename = fileName or os.path.basename(key)
+
+            # --------------------------------------------------
+            # Ensure extension from media_type
+            # --------------------------------------------------
+            if media_type:
+                ext = mimetypes.guess_extension(media_type) or ""
+                if ext and not filename.lower().endswith(ext):
+                    filename += ext
+
+            # --------------------------------------------------
+            # Final blob key
+            # --------------------------------------------------
+            blob_key = f"{key.rstrip('/')}/{filename}" if fileName else key
+
+            blob = self._bucket.blob(blob_key)
+
+            # --------------------------------------------------
+            # Metadata (must be string)
+            # --------------------------------------------------
+            safe_metadata = {k: str(v) for k, v in (metadata or {}).items()}
+            safe_metadata["filename"] = filename
+
+            blob.metadata = safe_metadata
+
+            # --------------------------------------------------
+            # Upload with content type
+            # --------------------------------------------------
+            blob.upload_from_string(
+                data,
+                content_type=media_type or "application/octet-stream",
+            )
+
+            # Persist metadata (required in GCS)
+            blob.patch()
+
+            self.logger.debug(f"GCS uploaded blob: {blob_key}, type={media_type}")
+
+            # --------------------------------------------------
+            # Return public URL
+            # --------------------------------------------------
+            return f"https://storage.googleapis.com/{self.bucket_name}/{blob_key}"
+
         except Exception as e:
             raise StorageError(f"GCS put failed: {str(e)}")
 
+    # ------------------------------------------------------------------
+    # Get object
+    # ------------------------------------------------------------------
+
     @retry(max_attempts=3, delay=1.0, exceptions=(StorageError,))
-    def get(self, key: str) -> bytes | None:
-        """Get object"""
+    def get(self, key: str) -> Optional[bytes]:
+        """Download object from GCS"""
         try:
-            if self._bucket:
-                blob = self._bucket.blob(key)
-                return blob.download_as_bytes()
+            if not self._bucket:
+                raise ConnectionError("GCS bucket not initialized")
+
+            blob = self._bucket.blob(key)
+
+            if not blob.exists():
+                return None
+
+            data = blob.download_as_bytes()
+
+            self.logger.debug(f"GCS downloaded blob: {key}")
+
+            return data
+
+        except NotFound:
             return None
+
         except Exception as e:
             raise StorageError(f"GCS get failed: {str(e)}")
 
+    # ------------------------------------------------------------------
+    # Delete object
+    # ------------------------------------------------------------------
+
     @retry(max_attempts=3, delay=1.0, exceptions=(StorageError,))
     def delete(self, key: str) -> bool:
-        """Delete object"""
+        """Delete object from GCS"""
         try:
-            if self._bucket:
-                blob = self._bucket.blob(key)
-                blob.delete()
-                return True
+            if not self._bucket:
+                raise ConnectionError("GCS bucket not initialized")
+
+            blob = self._bucket.blob(key)
+
+            if not blob.exists():
+                return False
+
+            blob.delete()
+
+            self.logger.debug(f"GCS deleted blob: {key}")
+
+            return True
+
+        except NotFound:
             return False
+
         except Exception as e:
             raise StorageError(f"GCS delete failed: {str(e)}")
 
+    # ------------------------------------------------------------------
+    # List objects
+    # ------------------------------------------------------------------
+
     @retry(max_attempts=3, delay=1.0, exceptions=(StorageError,))
     def list(self, prefix: str = "") -> List[str]:
-        """List objects with prefix"""
+        """List objects with optional prefix"""
         try:
-            if self._bucket:
-                blobs = self._bucket.list_blobs(prefix=prefix)
-                return [blob.name for blob in blobs]
-            return []
-        except Exception as e:
-            raise StorageError(f"GCS list failed: {str(e)}")
+            if not self._bucket:
+                raise ConnectionError("GCS bucket not initialized")
+
+            blobs = self._bucket.list_blobs(prefix=prefix)
+
+            results = [blob.name for blob in blobs]
 
+            self.logger.debug(f"GCS listed {len(results)} blobs (prefix={prefix})")
 
+            return results
+
+        except Exception as e:
+            raise StorageError(f"GCS list failed: {str(e)}")