altcodepro-polydb-python 2.5.4__py3-none-any.whl → 2.5.7__py3-none-any.whl

polydb/adapters/FirestoreAdapter.py

@@ -25,12 +25,21 @@ class FirestoreAdapter(NoSQLKVAdapter):
     Production-grade Firestore adapter with optional GCS overflow.
 
     Goals (matches your tests)
-    - Document id == pk (so querying {"id": ...} works)
+    - stored row keeps "id" == rk (the row key -- see NoSQLKVAdapter._get_pk_rk,
+      which derives rk from data.get(rk_field, ...) with rk_field defaulting
+      to "id"), so querying {"id": ...} matches the record's own id
     - patch() merges (preserves existing fields)
-    - delete() returns {"id": <pk>} and raises DatabaseError on missing
+    - delete() returns {"id": <rk>} and raises DatabaseError on missing
     - query_page() returns (rows, token) with stable pagination
     - Emulator support via FIRESTORE_EMULATOR_HOST
 
+    NOTE: the Firestore *document id* is still `pk` alone (see _doc_id) --
+    this adapter does not compose pk+rk into the physical document key the
+    way VercelKVAdapter/DynamoDBAdapter do. That is a pre-existing, separate
+    design constraint (two rows sharing the same pk collide onto one
+    document) left as-is here since fixing it is a bigger structural change
+    this pass didn't verify against a live Firestore/emulator.
+
     NOTE: google-cloud-firestore / google-cloud-storage are imported lazily
     inside the methods that use them, so installing them is only required when
     this adapter is actually used.
@@ -118,7 +127,7 @@ class FirestoreAdapter(NoSQLKVAdapter):
         return self._client.collection(self._collection_name(model))
 
     def _doc_id(self, pk: str) -> str:
-        # IMPORTANT for tests: doc_id == pk == row["id"]
+        # Document id is the partition key alone (not composed with rk).
         return str(pk)
 
     def _blob_key(self, model: type, pk: str, rk: str, checksum: str) -> str:
@@ -148,7 +157,7 @@ class FirestoreAdapter(NoSQLKVAdapter):
         blob.upload_from_string(data_bytes)
 
         ref: JsonDict = {
-            "id": pk,
+            "id": rk,
             "_pk": pk,
             "_rk": rk,
             "_overflow": True,
@@ -203,9 +212,15 @@ class FirestoreAdapter(NoSQLKVAdapter):
             collection = self._get_collection(model)
             doc_id = self._doc_id(pk)
 
-            # Ensure required identifiers exist for tests and convenience
+            # "id" is the row key (see NoSQLKVAdapter._get_pk_rk, which derives
+            # rk from data.get(rk_field, ...) with rk_field defaulting to "id").
+            # It is already present in `payload` via dict(data) above whenever
+            # the caller supplied one -- only fall back to `rk`, never `pk`,
+            # which would silently collapse every row's "id" to its partition
+            # key and break id-addressed lookups (the same corruption
+            # previously present here and in VercelKVAdapter/DynamoDBAdapter).
             payload: JsonDict = dict(data or {})
-            payload["id"] = pk
+            payload.setdefault("id", rk)
             payload["_pk"] = pk
             payload["_rk"] = rk
 
@@ -215,7 +230,8 @@ class FirestoreAdapter(NoSQLKVAdapter):
             else:
                 collection.document(doc_id).set(payload)
 
-            return {"id": pk}
+            # Return the full stored record, not just {"id": pk}.
+            return overflow_ref if overflow_ref is not None else payload
 
         except Exception as e:
             raise NoSQLError(f"Firestore put failed: {e}")
@@ -231,7 +247,7 @@ class FirestoreAdapter(NoSQLKVAdapter):
                 return None
 
             doc_data = snap.to_dict() or {}
-            doc_data.setdefault("id", pk)
+            doc_data.setdefault("id", rk)
 
             return self._resolve_overflow(doc_data)
 
@@ -275,8 +291,9 @@ class FirestoreAdapter(NoSQLKVAdapter):
             out: List[JsonDict] = []
             for d in docs:
                 row = d.to_dict() or {}
-                # Ensure id present for tests
-                row.setdefault("id", row.get("_pk") or d.id)
+                # Fall back to the row key, not the partition key, when "id"
+                # is absent from the stored document (legacy/incomplete data).
+                row.setdefault("id", row.get("_rk") or d.id)
                 out.append(self._resolve_overflow(row))
             return out
 
@@ -288,7 +305,7 @@ class FirestoreAdapter(NoSQLKVAdapter):
         """
         Test expectations:
         - deleting nonexistent raises sqlite3.DatabaseError
-        - delete returns {"id": pk}
+        - delete returns {"id": rk}
         - deletes overflow blob if present
         """
         try:
@@ -311,7 +328,7 @@ class FirestoreAdapter(NoSQLKVAdapter):
                         pass
 
             collection.document(doc_id).delete()
-            return {"id": pk}
+            return {"id": rk}
 
         except DatabaseError:
             raise
@@ -355,7 +372,7 @@ class FirestoreAdapter(NoSQLKVAdapter):
             rows: List[JsonDict] = []
             for d in docs:
                 row = d.to_dict() or {}
-                row.setdefault("id", row.get("_pk") or d.id)
+                row.setdefault("id", row.get("_rk") or d.id)
                 rows.append(self._resolve_overflow(row))
 
             next_token = None

polydb/adapters/GCPPubSubAdapter.py

@@ -237,3 +237,49 @@ class GCPPubSubAdapter(QueueAdapter):
 
         except Exception as e:
             raise QueueError(f"Pub/Sub ack failed: {e}")
+
+    # ---------------------------------------------------------
+    # extend -- real GCP Pub/Sub ModifyAckDeadline. delay()/cancel() are
+    # NOT overridden: a published Pub/Sub message is deliverable
+    # immediately (no "don't make this visible until N seconds from
+    # now" primitive exists), so the base class's NotImplementedError
+    # is the honest answer for both.
+    # ---------------------------------------------------------
+
+    # Pub/Sub's own real, documented maximum ack deadline. Google's API
+    # rejects anything above this with an invalid-argument error, so
+    # enforce it here rather than leaving callers to hit an opaque
+    # remote error.
+    MAX_ACK_DEADLINE_SECONDS = 600
+
+    @retry(max_attempts=3, delay=1.0, exceptions=(QueueError,))
+    def extend(
+        self, ack_id: str, queue_name: str = "default", *, visibility_timeout: int = 30
+    ) -> bool:
+        """Real Pub/Sub ModifyAckDeadline -- ack_id is the same ack_id
+        receive() returned for this message."""
+        if not ack_id:
+            raise QueueError("ack_id is required for Pub/Sub extend")
+        if visibility_timeout > self.MAX_ACK_DEADLINE_SECONDS:
+            raise QueueError(
+                f"Pub/Sub ack deadline cannot exceed {self.MAX_ACK_DEADLINE_SECONDS}s "
+                f"(GCP's own real maximum), got {visibility_timeout}"
+            )
+        try:
+            if not self._subscriber:
+                raise ConnectionError("Pub/Sub subscriber not initialized")
+
+            _, subscription = self._resolve_names(queue_name)
+            sub_path = self._subscription_path(subscription)
+
+            self._subscriber.modify_ack_deadline(
+                request={
+                    "subscription": sub_path,
+                    "ack_ids": [ack_id],
+                    "ack_deadline_seconds": visibility_timeout,
+                }
+            )
+            return True
+
+        except Exception as e:
+            raise QueueError(f"Pub/Sub extend failed: {e}")

polydb/adapters/GCPSecretManagerAdapter.py

@@ -0,0 +1,79 @@
+# src/polydb/adapters/GCPSecretManagerAdapter.py
+import os
+from typing import Optional
+
+from ..base.SecretsAdapter import SecretsAdapter
+from ..errors import ConnectionError
+from ..retry import retry
+
+
+class GCPSecretManagerAdapter(SecretsAdapter):
+    def __init__(self, project_id: str = ""):
+        super().__init__()
+
+        self.project_id = project_id or os.getenv("GOOGLE_CLOUD_PROJECT")
+        if not self.project_id:
+            raise ConnectionError("GOOGLE_CLOUD_PROJECT is not configured")
+
+        self._client = None
+        self._initialize_client()
+
+    def _initialize_client(self) -> None:
+        from google.cloud import secretmanager
+
+        self._client = secretmanager.SecretManagerServiceClient()
+
+    def _secret_path(self, key: str) -> str:
+        return f"projects/{self.project_id}/secrets/{key}"
+
+    @retry(max_attempts=3, delay=1.0, exceptions=(Exception,))
+    def get_secret(self, key: str) -> Optional[str]:
+        from google.api_core.exceptions import NotFound
+
+        try:
+            resp = self._client.access_secret_version(
+                name=f"{self._secret_path(key)}/versions/latest"
+            )
+        except NotFound:
+            return None
+        return resp.payload.data.decode("utf-8")
+
+    @retry(max_attempts=3, delay=1.0, exceptions=(Exception,))
+    def set_secret(self, key: str, value: str) -> None:
+        from google.api_core.exceptions import AlreadyExists
+
+        try:
+            self._client.create_secret(
+                request={
+                    "parent": f"projects/{self.project_id}",
+                    "secret_id": key,
+                    "secret": {"replication": {"automatic": {}}},
+                }
+            )
+        except AlreadyExists:
+            pass
+        self._client.add_secret_version(
+            request={
+                "parent": self._secret_path(key),
+                "payload": {"data": value.encode("utf-8")},
+            }
+        )
+
+    @retry(max_attempts=3, delay=1.0, exceptions=(Exception,))
+    def delete_secret(self, key: str) -> bool:
+        from google.api_core.exceptions import NotFound
+
+        try:
+            self._client.delete_secret(request={"name": self._secret_path(key)})
+            return True
+        except NotFound:
+            return False
+
+    @retry(max_attempts=3, delay=1.0, exceptions=(Exception,))
+    def list_secrets(self, prefix: str = "") -> list[str]:
+        secrets = self._client.list_secrets(request={"parent": f"projects/{self.project_id}"})
+        return [
+            s.name.rsplit("/", 1)[-1]
+            for s in secrets
+            if s.name.rsplit("/", 1)[-1].startswith(prefix)
+        ]

polydb/adapters/KafkaQueueAdapter.py

@@ -0,0 +1,332 @@
+# src/polydb/adapters/KafkaQueueAdapter.py
+import os
+import json
+import threading
+import uuid
+from typing import Any, Dict, List, Optional
+
+from ..base.QueueAdapter import QueueAdapter
+from ..errors import ConnectionError, QueueError
+from ..retry import retry
+from ..json_safe import json_safe
+
+
+class KafkaQueueAdapter(QueueAdapter):
+    """
+    Apache Kafka adapter using the synchronous `kafka-python` client, not
+    `aiokafka`. Every other adapter in this codebase (SQS, Azure Queue,
+    Pub/Sub, RabbitMQ, ...) is synchronous end to end -- pulling in an
+    asyncio-native client here would mean either spinning an event loop
+    per call or wrapping every call in asyncio.run(), both strictly worse
+    than using the client kafka-python already provides for exactly this
+    call shape.
+
+    queue_name IS the Kafka topic.
+    """
+
+    def __init__(
+        self,
+        bootstrap_servers: str = "",
+        group_id: str = "",
+        client_id: str = "",
+        security_protocol: str = "",
+        sasl_mechanism: str = "",
+        sasl_plain_username: str = "",
+        sasl_plain_password: str = "",
+        ssl_cafile: str = "",
+        auto_offset_reset: str = "earliest",
+    ):
+        super().__init__()
+
+        servers = bootstrap_servers or os.getenv("KAFKA_BOOTSTRAP_SERVERS", "localhost:9092")
+        self.bootstrap_servers = [s.strip() for s in servers.split(",") if s.strip()]
+
+        # A random per-instance consumer group by default so that two
+        # independently-constructed adapters (e.g. two test runs, or two
+        # unrelated callers in the same process) don't silently share
+        # partition assignment / offsets with each other. Callers that
+        # actually want shared work-queue semantics across processes pass
+        # an explicit group_id (or set KAFKA_GROUP_ID).
+        self.group_id = (
+            group_id or os.getenv("KAFKA_GROUP_ID") or f"polydb-{uuid.uuid4().hex[:12]}"
+        )
+        self.client_id = client_id or os.getenv("KAFKA_CLIENT_ID", "polydb")
+        self.auto_offset_reset = auto_offset_reset
+
+        # Auth/TLS are all optional -- PLAINTEXT (no auth) is the default
+        # so local/dev brokers work with zero extra config, matching how
+        # e.g. SQSAdapter defaults endpoint_url to "" (real AWS) rather
+        # than requiring LocalStack settings to be supplied.
+        self.security_protocol = security_protocol or os.getenv(
+            "KAFKA_SECURITY_PROTOCOL", "PLAINTEXT"
+        )
+        self.sasl_mechanism = sasl_mechanism or os.getenv("KAFKA_SASL_MECHANISM") or None
+        self.sasl_plain_username = (
+            sasl_plain_username or os.getenv("KAFKA_SASL_USERNAME") or None
+        )
+        self.sasl_plain_password = (
+            sasl_plain_password or os.getenv("KAFKA_SASL_PASSWORD") or None
+        )
+        self.ssl_cafile = ssl_cafile or os.getenv("KAFKA_SSL_CAFILE") or None
+
+        self._producer: Any = None
+        self._consumers: Dict[str, Any] = {}  # topic -> KafkaConsumer
+
+        # message_id -> (TopicPartition, offset_to_commit). Populated by
+        # receive(), consumed (and popped) by ack()/delete(). See the long
+        # comment on receive() for why committing only happens here.
+        self._pending: Dict[str, Any] = {}
+
+        self._lock = threading.Lock()
+
+    # ---------------------------------------------------------
+    # Client initialization
+    # ---------------------------------------------------------
+
+    def _client_kwargs(self) -> Dict[str, Any]:
+        kwargs: Dict[str, Any] = {
+            "bootstrap_servers": self.bootstrap_servers,
+            "security_protocol": self.security_protocol,
+        }
+        if self.sasl_mechanism:
+            kwargs["sasl_mechanism"] = self.sasl_mechanism
+            kwargs["sasl_plain_username"] = self.sasl_plain_username
+            kwargs["sasl_plain_password"] = self.sasl_plain_password
+        if self.ssl_cafile:
+            kwargs["ssl_cafile"] = self.ssl_cafile
+        return kwargs
+
+    def _get_producer(self):
+        from kafka import KafkaProducer
+
+        if self._producer is not None:
+            return self._producer
+
+        with self._lock:
+            if self._producer is not None:
+                return self._producer
+            try:
+                self._producer = KafkaProducer(
+                    client_id=self.client_id,
+                    # We JSON-encode to bytes ourselves (matches json_safe
+                    # usage elsewhere in this codebase), so the serializer
+                    # is a passthrough rather than kafka-python's own.
+                    value_serializer=lambda v: v,
+                    **self._client_kwargs(),
+                )
+                self.logger.info(
+                    f"Initialized Kafka producer (bootstrap={self.bootstrap_servers})"
+                )
+            except Exception as e:
+                raise ConnectionError(f"Kafka producer init failed: {e}")
+            return self._producer
+
+    def _get_consumer(self, topic: str):
+        if topic in self._consumers:
+            return self._consumers[topic]
+
+        from kafka import KafkaConsumer
+
+        with self._lock:
+            if topic in self._consumers:
+                return self._consumers[topic]
+            try:
+                consumer = KafkaConsumer(
+                    topic,
+                    group_id=self.group_id,
+                    client_id=self.client_id,
+                    # Manual commits only -- see receive()'s docstring for
+                    # why offsets are committed exclusively from ack()/
+                    # delete(), never automatically here.
+                    enable_auto_commit=False,
+                    auto_offset_reset=self.auto_offset_reset,
+                    **self._client_kwargs(),
+                )
+                self.logger.info(f"Initialized Kafka consumer (topic={topic}, group={self.group_id})")
+            except Exception as e:
+                raise ConnectionError(f"Kafka consumer init failed: {e}")
+            self._consumers[topic] = consumer
+            return consumer
+
+    # ---------------------------------------------------------
+    # Queue operations
+    # ---------------------------------------------------------
+
+    @retry(max_attempts=3, delay=1.0, exceptions=(QueueError,))
+    def send(self, message: Dict[str, Any], queue_name: str = "default") -> str:
+        """Produce a message to `queue_name` (the Kafka topic)."""
+        try:
+            producer = self._get_producer()
+            body = (
+                json.dumps(message, default=json_safe).encode("utf-8")
+                if not isinstance(message, (bytes, bytearray))
+                else message
+            )
+
+            future = producer.send(queue_name, value=body)
+            # kafka-python's send() is async by default (it returns a
+            # FutureRecordMetadata immediately, before the broker has
+            # necessarily accepted the record). .get() blocks for that ack
+            # so send() returns only once the message is durably produced
+            # -- matching every other adapter's synchronous
+            # send()->message_id contract instead of firing-and-forgetting.
+            record = future.get(timeout=10)
+            return f"{record.partition}-{record.offset}"
+
+        except Exception as e:
+            raise QueueError(f"Kafka 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]]:
+        """
+        Poll up to `max_messages` from `queue_name`'s consumer group.
+
+        Deliberately does NOT commit offsets here (the consumer is created
+        with enable_auto_commit=False too). A message only becomes "done"
+        from Kafka's point of view once ack()/delete() commits its offset.
+        If the caller crashes after receive() but before ack(), nothing
+        was ever committed, so the next poll -- this process restarted, or
+        any other consumer sharing this group_id -- redelivers the message
+        from the same offset. That's the same at-least-once shape SQS's
+        visibility timeout, Azure's visibility timeout, and Pub/Sub's
+        unacked-redelivery already give this codebase's other queue
+        adapters. Auto-committing inside receive() would instead mark a
+        message "consumed" the instant it's handed out, which loses
+        redelivery on a crash mid-processing -- effectively at-most-once,
+        not the at-least-once contract the rest of this codebase relies on
+        (WorkerPool retries on a failed/never-acked message).
+        """
+        try:
+            from kafka import TopicPartition
+
+            consumer = self._get_consumer(queue_name)
+
+            out: List[Dict[str, Any]] = []
+            # poll() returns whatever's ready in a single batch, which may
+            # be less than max_messages even when more exists -- loop
+            # (bounded) rather than assuming one poll() satisfies the ask.
+            attempts = 0
+            while len(out) < max_messages and attempts < 5:
+                attempts += 1
+                remaining = max_messages - len(out)
+                batches = consumer.poll(timeout_ms=1000, max_records=remaining)
+                if not batches:
+                    break
+
+                for tp, records in batches.items():
+                    for record in records:
+                        message_id = f"{record.partition}-{record.offset}"
+
+                        try:
+                            body = json.loads(record.value.decode("utf-8"))
+                        except Exception:
+                            body = record.value.decode("utf-8", errors="replace")
+
+                        out.append(
+                            {
+                                "id": message_id,
+                                "receipt_handle": message_id,
+                                "body": body,
+                                "topic": tp.topic,
+                                "partition": tp.partition,
+                                "offset": record.offset,
+                            }
+                        )
+
+                        # Kafka commit semantics: the committed offset is
+                        # "the next record to read", so we store
+                        # offset + 1, not offset itself.
+                        self._pending[message_id] = (
+                            TopicPartition(tp.topic, tp.partition),
+                            record.offset + 1,
+                        )
+
+                        if len(out) >= max_messages:
+                            break
+                    if len(out) >= max_messages:
+                        break
+
+            return out
+
+        except Exception as e:
+            raise QueueError(f"Kafka receive failed: {e}")
+
+    def _commit(self, message_id: str, queue_name: str) -> bool:
+        from kafka import OffsetAndMetadata
+
+        pending = self._pending.pop(message_id, None)
+        if pending is None:
+            # Unknown or already-committed id -- treat as a no-op success.
+            # Matches VercelQueueAdapter/BlockchainQueueAdapter's existing
+            # convention of a redundant ack being harmless rather than an
+            # error.
+            return False
+
+        tp, next_offset = pending
+        consumer = self._consumers.get(queue_name)
+        if consumer is None:
+            raise QueueError(
+                f"No active Kafka consumer for topic '{queue_name}' to commit offset against"
+            )
+
+        try:
+            consumer.commit({tp: OffsetAndMetadata(next_offset, None)})
+            return True
+        except Exception as e:
+            raise QueueError(f"Kafka offset commit failed: {e}")
+
+    def delete(self, message_id: str, queue_name: str = "default", pop_receipt: str = "") -> bool:
+        """
+        Delete == commit the offset. Kafka has no notion of deleting a
+        single record independent of consumer offsets, so -- matching the
+        majority convention among this codebase's other adapters (SQS,
+        Pub/Sub: ack is delete) -- delete() and ack() do the same thing.
+        """
+        return self._commit(message_id, queue_name)
+
+    def ack(self, ack_id: str, queue_name: str = "default") -> bool:
+        """Explicit ACK: commits the consumed offset for `ack_id`."""
+        if not ack_id:
+            raise QueueError("ack_id is required for Kafka ack")
+        return self._commit(ack_id, queue_name)
+
+    def nack(self, ack_id: str, queue_name: str = "default") -> bool:
+        """
+        Documented no-op, not a design gap: receive()'s own docstring
+        already establishes that this adapter deliberately never commits
+        an offset until ack()/delete() does. That means a message is
+        already effectively "nacked" -- redeliverable to this (or any
+        other) consumer sharing group_id -- the instant it's received
+        and not yet acked; there is no separate broker-side "put it back"
+        call the way AMQP's basic_nack is, because nothing was ever
+        marked done in the first place.
+
+        This only pops the pending entry (mirroring _commit's own
+        "unknown/already-handled id is a harmless no-op" convention, and
+        returning the same True-if-there-was-something-to-act-on /
+        False-if-not shape _commit does) so a caller that explicitly
+        nacks doesn't also get to ack() the same id afterward -- it
+        exists for API-shape consistency with the other adapters'
+        nack(), not because Kafka needs an explicit call here to achieve
+        the redelivery.
+        """
+        return self._pending.pop(ack_id, None) is not None
+
+    def close(self) -> None:
+        """
+        Flush the producer and close every consumer. Not part of
+        QueueAdapter's abstract contract (none of the other adapters need
+        it -- boto3/Azure SDK/Pub/Sub clients don't hold a persistent
+        local socket the way a Kafka producer/consumer does), but good
+        hygiene for a long-lived adapter instance to call explicitly.
+        """
+        if self._producer is not None:
+            try:
+                self._producer.close()
+            except Exception:
+                pass
+        for consumer in self._consumers.values():
+            try:
+                consumer.close()
+            except Exception:
+                pass