Application: accept "quix_portal_api" param (#902)

* Application: accept "quix_portal_api" param
* Concentrate env variable access in QuixEnvironment

Author: daniil-quix

quixstreams/app.py

@@ -1,22 +1,12 @@
 import contextlib
 import functools
 import logging
-import os
 import signal
 import time
 import warnings
 from collections import defaultdict
 from pathlib import Path
-from typing import (
-    Callable,
-    List,
-    Literal,
-    Optional,
-    Protocol,
-    Tuple,
-    Type,
-    Union,
-)
+from typing import Callable, List, Literal, Optional, Protocol, Tuple, Type, Union
 
 from confluent_kafka import TopicPartition
 from pydantic import AliasGenerator, Field
@@ -45,12 +35,14 @@
     TopicManager,
 )
 from .platforms.quix import (
+    DEFAULT_PORTAL_API_URL,
     QuixKafkaConfigsBuilder,
     QuixTopicManager,
     check_state_dir,
     check_state_management_enabled,
     is_quix_deployment,
 )
+from .platforms.quix.env import QUIX_ENVIRONMENT
 from .processing import ProcessingContext
 from .runtracker import RunTracker
 from .sinks import SinkManager
@@ -134,6 +126,7 @@ def __init__(
         broker_address: Optional[Union[str, ConnectionConfig]] = None,
         *,
         quix_sdk_token: Optional[str] = None,
+        quix_portal_api: Optional[str] = None,
         consumer_group: Optional[str] = None,
         auto_offset_reset: AutoOffsetReset = "latest",
         commit_interval: float = 5.0,
@@ -172,6 +165,11 @@ def __init__(
             Linked Environment Variable: `Quix__Sdk__Token`.
             Default: None (if not run on Quix Cloud)
               >***NOTE:*** the environment variable is set for you in the Quix Cloud
+        :param quix_portal_api: If using the Quix Cloud, the cluster API URL to use.
+            Use it to connect to the dedicated Quix Cloud environment.
+            Linked Environment Variable: `Quix__Portal__Api`.
+            Default: `https://portal-api.platform.quix.io/`.
+              >***NOTE:*** the environment variable is set for you in the Quix Cloud
         :param consumer_group: Kafka consumer group.
             Passed as `group.id` to `confluent_kafka.Consumer`.
             Linked Environment Variable: `Quix__Consumer_Group`.
@@ -244,19 +242,20 @@ def __init__(
         consumer_extra_config = consumer_extra_config or {}
 
         if state_dir is None:
-            state_dir = os.getenv(
-                "Quix__State__Dir", "/app/state" if is_quix_deployment() else "state"
+            state_dir = QUIX_ENVIRONMENT.state_dir or (
+                "/app/state" if is_quix_deployment() else "state"
             )
         state_dir = Path(state_dir)
 
-        # We can't use os.getenv as defaults (and have testing work nicely)
-        # since it evaluates getenv when the function is defined.
-        # In general this is just a most robust approach.
-        broker_address = broker_address or os.getenv("Quix__Broker__Address")
-        quix_sdk_token = quix_sdk_token or os.getenv("Quix__Sdk__Token")
+        broker_address = broker_address or QUIX_ENVIRONMENT.broker_address
+        quix_sdk_token = quix_sdk_token or QUIX_ENVIRONMENT.sdk_token
+        quix_portal_api = (
+            quix_portal_api or QUIX_ENVIRONMENT.portal_api or DEFAULT_PORTAL_API_URL
+        )
 
-        if not consumer_group:
-            consumer_group = os.getenv("Quix__Consumer_Group", "quixstreams-default")
+        consumer_group = (
+            consumer_group or QUIX_ENVIRONMENT.consumer_group or "quixstreams-default"
+        )
 
         if broker_address:
             # If broker_address is passed to the app it takes priority over any quix configuration
@@ -276,8 +275,8 @@ def __init__(
                     )
             elif quix_sdk_token:
                 quix_app_source = "Quix SDK Token"
-                quix_config_builder = QuixKafkaConfigsBuilder(
-                    quix_sdk_token=quix_sdk_token
+                quix_config_builder = QuixKafkaConfigsBuilder.from_credentials(
+                    quix_sdk_token=quix_sdk_token, quix_portal_api=quix_portal_api
                 )
             else:
                 raise ValueError(
@@ -287,7 +286,8 @@ def __init__(
             # SDK Token or QuixKafkaConfigsBuilder were provided
             logger.info(
                 f"{quix_app_source} detected; "
-                f"the application will connect to Quix Cloud brokers"
+                f"the application will connect to Quix Cloud brokers "
+                f'(quix_portal_api="{quix_portal_api}")'
             )
             self._topic_manager_factory = functools.partial(
                 QuixTopicManager, quix_config_builder=quix_config_builder

quixstreams/platforms/quix/api.py

@@ -14,8 +14,6 @@
 
 __all__ = ("QuixPortalApiService",)
 
-DEFAULT_PORTAL_API_URL = "https://portal-api.platform.quix.io/"
-
 
 class QuixPortalApiService:
     """
@@ -34,13 +32,11 @@ class QuixPortalApiService:
     def __init__(
         self,
         auth_token: str,
-        portal_api: Optional[str] = None,
+        portal_api: str,
         api_version: Optional[str] = None,
         default_workspace_id: Optional[str] = None,
     ):
-        self._portal_api_url = (
-            portal_api or QUIX_ENVIRONMENT.portal_api or DEFAULT_PORTAL_API_URL
-        )
+        self._portal_api_url = portal_api
         if not auth_token:
             raise MissingConnectionRequirements(
                 f"A Quix Cloud auth token (SDK or PAT) is required; "

quixstreams/platforms/quix/config.py

@@ -20,10 +20,12 @@
     UndefinedQuixWorkspaceId,
 )
 
+__all__ = ("QuixKafkaConfigsBuilder", "QuixApplicationConfig", "DEFAULT_PORTAL_API_URL")
+
 logger = logging.getLogger(__name__)
 
-__all__ = ("QuixKafkaConfigsBuilder", "QuixApplicationConfig")
 
+DEFAULT_PORTAL_API_URL = "https://portal-api.platform.quix.io/"
 QUIX_CONNECTIONS_MAX_IDLE_MS = 3 * 60 * 1000
 QUIX_METADATA_MAX_AGE_MS = 3 * 60 * 1000
 
@@ -95,32 +97,21 @@ class QuixKafkaConfigsBuilder:
     It also currently handles the app_auto_create_topics setting for Quix Applications.
     """
 
-    # TODO: Consider a workspace class?
     def __init__(
         self,
-        quix_sdk_token: Optional[str] = None,
+        quix_portal_api_service: QuixPortalApiService,
         workspace_id: Optional[str] = None,
-        quix_portal_api_service: Optional[QuixPortalApiService] = None,
         timeout: float = 30,
         topic_create_timeout: float = 60,
     ):
         """
-        :param quix_portal_api_service: A QuixPortalApiService instance (else generated)
+        :param quix_portal_api_service: A QuixPortalApiService instance
         :param workspace_id: A valid Quix Workspace ID (else searched for)
         """
-        if quix_sdk_token:
-            self.api = QuixPortalApiService(
-                default_workspace_id=workspace_id, auth_token=quix_sdk_token
-            )
-        elif quix_portal_api_service:
-            self.api = quix_portal_api_service
-        else:
-            raise ValueError(
-                'Either "quix_sdk_token" or "quix_portal_api_service" must be provided'
-            )
 
+        self._api = quix_portal_api_service
         try:
-            self._workspace_id = workspace_id or self.api.default_workspace_id
+            self._workspace_id = workspace_id or self._api.default_workspace_id
         except UndefinedQuixWorkspaceId:
             self._workspace_id = ""
             logger.warning(
@@ -136,6 +127,30 @@ def __init__(
         self._timeout = timeout
         self._topic_create_timeout = topic_create_timeout
 
+    @classmethod
+    def from_credentials(
+        cls,
+        quix_sdk_token: str,
+        quix_portal_api: str = DEFAULT_PORTAL_API_URL,
+        workspace_id: Optional[str] = None,
+        timeout: float = 30,
+        topic_create_timeout: float = 60,
+    ) -> "QuixKafkaConfigsBuilder":
+        """
+        Initialize class using the quix_sdk_token and quix_portal_api params.
+        """
+        api_service = QuixPortalApiService(
+            default_workspace_id=workspace_id,
+            auth_token=quix_sdk_token,
+            portal_api=quix_portal_api,
+        )
+        return cls(
+            quix_portal_api_service=api_service,
+            workspace_id=workspace_id,
+            timeout=timeout,
+            topic_create_timeout=topic_create_timeout,
+        )
+
     @property
     def workspace_id(self) -> str:
         if not self._workspace_id:
@@ -250,12 +265,12 @@ def search_for_workspace(
         if not workspace_name_or_id:
             workspace_name_or_id = self._workspace_id
         try:
-            return self.api.get_workspace(
+            return self._api.get_workspace(
                 workspace_id=workspace_name_or_id, timeout=timeout
             )
         except HTTPError:
             # check to see if they provided the workspace name instead
-            ws_list = self.api.get_workspaces(timeout=timeout)
+            ws_list = self._api.get_workspaces(timeout=timeout)
             for ws in ws_list:
                 if ws["name"] == workspace_name_or_id:
                     return ws
@@ -315,7 +330,7 @@ def search_workspace_for_topic(
 
         :return: the workspace_id if success, else None
         """
-        topics = self.api.get_topics(
+        topics = self._api.get_topics(
             workspace_id=workspace_id,
             timeout=timeout if timeout is not None else self._timeout,
         )
@@ -338,7 +353,7 @@ def search_for_topic_workspace(
 
         :return: workspace data dict if topic search success, else None
         """
-        ws_list = self.api.get_workspaces(
+        ws_list = self._api.get_workspaces(
             timeout=timeout if timeout is not None else self._timeout
         )
         if len(ws_list) == 1:
@@ -376,7 +391,7 @@ def create_topic(self, topic: Topic, timeout: Optional[float] = None) -> dict:
             f'Creating topic "{topic.name}" '
             f'with a config: "{topic.create_config.as_dict() if topic.create_config is not None else {} }"'
         )
-        resp = self.api.post_topic(
+        resp = self._api.post_topic(
             topic_name=topic.name,
             workspace_id=self.workspace_id,
             topic_partitions=cfg.num_partitions,
@@ -473,14 +488,14 @@ def get_topic(self, topic: Topic, timeout: Optional[float] = None) -> dict:
         :return: response dict of the topic info if topic found, else None
         :raises QuixApiRequestFailure: when topic does not exist
         """
-        return self.api.get_topic(
+        return self._api.get_topic(
             topic_name=topic.name,
             workspace_id=self.workspace_id,
             timeout=timeout if timeout is not None else self._timeout,
         )
 
     def get_topics(self, timeout: Optional[float] = None) -> List[dict]:
-        return self.api.get_topics(
+        return self._api.get_topics(
             workspace_id=self.workspace_id,
             timeout=timeout if timeout is not None else self._timeout,
         )
@@ -490,7 +505,7 @@ def _get_librdkafka_connection_config(self) -> ConnectionConfig:
         Get the full client config required to authenticate a confluent-kafka
         client to a Quix platform broker/workspace as a ConnectionConfig
         """
-        librdkafka_dict = self.api.get_librdkafka_connection_config(self.workspace_id)
+        librdkafka_dict = self._api.get_librdkafka_connection_config(self.workspace_id)
         if (cert := librdkafka_dict.pop("ssl.ca.cert", None)) is not None:
             librdkafka_dict["ssl.ca.pem"] = base64.b64decode(cert)
         return ConnectionConfig.from_librdkafka_dict(librdkafka_dict)

quixstreams/platforms/quix/env.py

@@ -6,14 +6,17 @@
 
 class QuixEnvironment:
     """
-    Class to access various Quix platform environment settings
+    A class to access various Quix Streams environment variables
     """
 
     SDK_TOKEN = "Quix__Sdk__Token"  # noqa: S105
+    BROKER_ADDRESS = "Quix__Broker__Address"
     PORTAL_API = "Quix__Portal__Api"
     WORKSPACE_ID = "Quix__Workspace__Id"
     DEPLOYMENT_ID = "Quix__Deployment__Id"
     STATE_MANAGEMENT_ENABLED = "Quix__Deployment__State__Enabled"
+    STATE_DIR = "Quix__State__Dir"
+    CONSUMER_GROUP = "Quix__Consumer_Group"
 
     @property
     def state_management_enabled(self) -> bool:
@@ -44,22 +47,40 @@ def workspace_id(self) -> Optional[str]:
         return os.environ.get(self.WORKSPACE_ID)
 
     @property
-    def portal_api(self) -> Optional[str]:
+    def state_dir(self) -> Optional[str]:
         """
-        Return Quix Portal API url if set
+        Return application state directory on Quix.
+        :return: path to state dir
+        """
+        return os.environ.get(self.STATE_DIR)
 
-        :return: portal API URL or None
+    @property
+    def portal_api(self) -> Optional[str]:
+        """
+        Quix Portal API URL
         """
         return os.environ.get(self.PORTAL_API)
 
     @property
-    def state_dir(self) -> str:
+    def broker_address(self) -> Optional[str]:
         """
-        Return application state directory on Quix.
-        :return: path to state dir
+        Kafka broker address
+        """
+        return os.environ.get(self.BROKER_ADDRESS)
+
+    @property
+    def sdk_token(self) -> Optional[str]:
+        """
+        Quix SDK token
+        """
+        return os.environ.get(self.SDK_TOKEN)
+
+    @property
+    def consumer_group(self) -> Optional[str]:
+        """
+        Kafka consumer group
         """
-        # TODO: Use env variables instead when they're available
-        return "/app/state"
+        return os.environ.get(self.CONSUMER_GROUP)
 
 
 QUIX_ENVIRONMENT = QuixEnvironment()

quixstreams/platforms/quix/exceptions.py

@@ -6,6 +6,10 @@
     "MissingConnectionRequirements",
     "UndefinedQuixWorkspaceId",
     "QuixApiRequestFailure",
+    "MultipleWorkspaces",
+    "NoWorkspaceFound",
+    "QuixCreateTopicFailure",
+    "QuixCreateTopicTimeout",
 )
 
 

quixstreams/sources/core/kafka/quix.py

@@ -4,8 +4,7 @@
 from quixstreams.kafka import AutoOffsetReset
 from quixstreams.models.serializers import DeserializerType
 from quixstreams.models.topics import Topic
-from quixstreams.platforms.quix import QuixKafkaConfigsBuilder
-from quixstreams.platforms.quix.api import QuixPortalApiService
+from quixstreams.platforms.quix import DEFAULT_PORTAL_API_URL, QuixKafkaConfigsBuilder
 from quixstreams.sources import (
     ClientConnectFailureCallback,
     ClientConnectSuccessCallback,
@@ -83,12 +82,10 @@ def __init__(
 
         self._short_topic = topic
         self._quix_workspace_id = quix_workspace_id
-        self._quix_config = QuixKafkaConfigsBuilder(
-            quix_portal_api_service=QuixPortalApiService(
-                default_workspace_id=quix_workspace_id,
-                auth_token=quix_sdk_token,
-                portal_api=quix_portal_api,
-            )
+        self._quix_config = QuixKafkaConfigsBuilder.from_credentials(
+            quix_sdk_token=quix_sdk_token,
+            quix_portal_api=quix_portal_api or DEFAULT_PORTAL_API_URL,
+            workspace_id=quix_workspace_id,
         )
 
         quix_topic = self._quix_config.convert_topic_response(