chalkpy 2.98.6__py3-none-any.whl → 2.99.1__py3-none-any.whl

chalk/_version.py

@@ -1 +1 @@
-__version__ = "2.98.6"
+__version__ = "2.99.1"

chalk/client/client_impl.py

@@ -1033,6 +1033,10 @@ class ChalkAPIClientImpl(ChalkClient):
         self._branch: str | None = branch
         self._skip_token_cache: bool = _skip_cache
         self._api_server: str | None = api_server
+        self._env_id_to_engine_url_map: Mapping[str, str] | None = None
+        self._env_id_to_env_name_map: Mapping[str, str] | None = None
+        self._env_name_to_env_id_map: Mapping[str, str] | None = None
+        self._access_token: str | None = None
 
         self._default_headers = {
             "Accept": "application/json",
@@ -1043,7 +1047,7 @@ class ChalkAPIClientImpl(ChalkClient):
         if additional_headers:
             self._default_headers.update(additional_headers)
 
-        self._primary_environment = environment
+        self._primary_environment: EnvironmentId | None = environment
 
         self.__class__.latest_client = self
         if notebook.is_notebook():
@@ -1116,11 +1120,21 @@ class ChalkAPIClientImpl(ChalkClient):
             creds = ExchangeCredentialsResponse(**response_json)
         except ValidationError:
             raise HTTPError(response=resp)
-        self._default_headers["Authorization"] = f"Bearer {creds.access_token}"
+
+        self._save_credentials_response(creds)
+
         # FIXME: We should NOT be using the X-Chalk-Client-Id for anything, as it is NOT authenticated
         self._default_headers["X-Chalk-Client-Id"] = self._client_id
+
+    def _save_credentials_response(self, creds: ExchangeCredentialsResponse):
+        self._access_token = creds.access_token
+        self._default_headers["Authorization"] = f"Bearer {creds.access_token}"
         if self._primary_environment is None:
             self._primary_environment = creds.primary_environment
+        self._env_id_to_engine_url_map = creds.engines
+        self._env_id_to_env_name_map = creds.environment_id_to_name
+        if self._env_id_to_env_name_map is not None:
+            self._env_name_to_env_id_map = {v: k for k, v in self._env_id_to_env_name_map.items()}
 
     def _get_headers(
         self,
@@ -1383,6 +1397,53 @@ https://docs.chalk.ai/docs/debugging-queries#resolver-replay
             method=method, headers=headers, url=url, json=json_body, data=data, timeout=timeout_value
         )
 
+    def _get_engine_host(self, environment_override: Optional[str]) -> str | None:
+        """
+        Returns the host to use for data-plane requests to the engine. May fall back to
+        the metadata plane api server host if no engine host can be determined.
+        :param environment_override: Optional user-provided environment name or id.
+        :return: Host for direct engine (data-plane) requests
+        """
+
+        # always respect user-provided query_server override
+        if self._query_server is not None:
+            return self._query_server
+
+        # if we have no environment information, we likely have not done creds exchange yet.
+        # there's no way to determine an engine host here, so just fall back to api server
+        env_id_or_name: str | None = environment_override or self._primary_environment
+        if env_id_or_name is None:
+            return self._api_server
+
+        # get the correct engine url given a valid env id or env name
+        if self._env_id_to_engine_url_map is not None:
+            if env_id_or_name in self._env_id_to_engine_url_map:
+                return self._env_id_to_engine_url_map[env_id_or_name]
+            elif self._env_name_to_env_id_map is not None and env_id_or_name in self._env_name_to_env_id_map:
+                env_id = self._env_name_to_env_id_map[env_id_or_name]
+                return self._env_id_to_engine_url_map.get(env_id, self._api_server)
+
+        # env name or id is invalid, or we failed to auth, so fall back to api server
+        return self._api_server
+
+    def _get_request_host(self, metadata_request: bool, environment_override: Optional[str]) -> str:
+        """
+        Returns the hostname to use for any requests made by the Chalk Client, whether they are
+        metadata-plane-only requests or data-plane requests (to the engine).
+        :param metadata_request: Is this a metadata-plane-only request? If so, we always use the API server.
+        :param environment_override: Optional user-provided environment name or id.
+        :return: Host to use for any request made by the Chalk Client
+        """
+        if metadata_request:
+            host = self._api_server
+        else:
+            host = self._get_engine_host(environment_override)
+
+        return host or "https://api.chalk.ai"
+
+    def _has_authorization_header(self) -> bool:
+        return "Authorization" in self._default_headers or "authorization" in self._default_headers
+
     def _request(
         self,
         method: str,
@@ -1399,25 +1460,18 @@ https://docs.chalk.ai/docs/debugging-queries#resolver-replay
         timeout: float | None | ellipsis = ...,
         connect_timeout: float | None | ellipsis = ...,
     ) -> T | requests.Response:
+        # If we don't have an access token, we cannot make a request.
+        # Always fetch an access token if we do not have one before making a request.
+        # If we fetched an access token during this call and we get 401/403 on this request,
+        # then do not retry the credential exchange again.
+        # Additional Note: if the user has explicitly provided an Authorization header, we should
+        # not attempt to exchange credentials, as they are explicitly managing auth themselves.
         allow_credential_exchange: bool = True
-
-        if metadata_request or self._query_server is None:
-            host = self._api_server
-        else:
-            host = self._query_server
-
-        if host is None:
-            # We definitively need to exchange credentials to get a host
+        if self._access_token is None and not self._has_authorization_header():
             self._exchange_credentials()
             allow_credential_exchange = False
 
-            # After exchanging credentials, the api server is never none
-            assert self._api_server is not None
-
-            if metadata_request or self._query_server is None:
-                host = self._api_server
-            else:
-                host = self._query_server
+        host = self._get_request_host(metadata_request, environment_override)
 
         r = self._do_request_inner(
             method=method,
@@ -1438,14 +1492,7 @@ https://docs.chalk.ai/docs/debugging-queries#resolver-replay
             # It is possible that credentials expired, or that we changed permissions since we last
             # got a token. Exchange them and try again
             self._exchange_credentials()
-
-            # After exchanging credentials, the api server is never null
-            assert self._api_server is not None
-
-            if metadata_request or self._query_server is None:
-                host = self._api_server
-            else:
-                host = self._query_server
+            host = self._get_request_host(metadata_request, environment_override)
 
             r = self._do_request_inner(
                 method=method,
@@ -3987,9 +4034,8 @@ https://docs.chalk.ai/cli/apply
             environment_override=context.environment,
             preview_deployment_id=preview_deployment_id,
             branch=branch,
-            # If using multiple computers, then we must route through the metadata server
-            # So we can actually spin up multiple pods
-            metadata_request=request.use_multiple_computers,
+            # all offline query requests should be routed through api server
+            metadata_request=True,
         )
         return response
 
@@ -4065,19 +4111,6 @@ https://docs.chalk.ai/cli/apply
             branch=None,
         )
 
-    def _get_query_inputs(
-        self, job_id: uuid.UUID, environment: Optional[EnvironmentId], branch: Optional[BranchId]
-    ) -> GetOfflineQueryJobResponse:
-        return self._request(
-            method="GET",
-            uri=f"/v2/offline_query_inputs/{job_id}",
-            response=GetOfflineQueryJobResponse,
-            environment_override=environment,
-            json=None,
-            preview_deployment_id=None,
-            branch=branch,
-        )
-
     def _get_dataset_from_name_or_id(
         self,
         *,

chalk/client/models.py

@@ -454,6 +454,8 @@ class ExchangeCredentialsResponse(BaseModel):
     api_server: str
     primary_environment: Optional[str] = None
     engines: Optional[Mapping[str, str]] = None
+    grpc_engines: Optional[Mapping[str, str]] = None
+    environment_id_to_name: Optional[Mapping[str, str]] = None
 
 
 class OfflineQueryInput(BaseModel):

chalk/features/resolver.py

@@ -4541,7 +4541,14 @@ def validate_underscore_expression(
     name: str,
 ) -> None:
     """Main entry point for validating a single underscore expression."""
-    from chalk.features.underscore import Underscore, UnderscoreAttr, UnderscoreCall, UnderscoreCast, UnderscoreFunction
+    from chalk.features.underscore import (
+        Underscore,
+        UnderscoreAttr,
+        UnderscoreCall,
+        UnderscoreCast,
+        UnderscoreFunction,
+        UnderscoreRoot,
+    )
 
     # First validate it's an underscore
     if not isinstance(expression, Underscore):  # pyright: ignore[reportUnnecessaryIsInstance]
@@ -4562,6 +4569,8 @@ def validate_underscore_expression(
         validate_underscore_call(expression, message_type, error_builder, name)
     elif isinstance(expression, UnderscoreCast):
         validate_underscore_cast(expression, message_type, error_builder, name)
+    elif isinstance(expression, UnderscoreRoot):
+        pass
     else:
         # Catch-all for any other underscore types
         error_builder.add_diagnostic(

{chalkpy-2.98.6.dist-info → chalkpy-2.99.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: chalkpy
-Version: 2.98.6
+Version: 2.99.1
 Summary: Python SDK for Chalk
 Author: Chalk AI, Inc.
 Project-URL: Homepage, https://chalk.ai

{chalkpy-2.98.6.dist-info → chalkpy-2.99.1.dist-info}/RECORD

@@ -1,5 +1,5 @@
 chalk/__init__.py,sha256=vKsx9-cl5kImlVWGHVRYO6bweBm79NAzGs3l36u71wM,2657
-chalk/_version.py,sha256=lCU8-hElWBDbqghv9vMcFvFH7YnIyamLkmO8KLLgp4g,23
+chalk/_version.py,sha256=veq7UDoQd0z2ACc6n5r66tpMpzpZd-wj3jbKxjfT1KQ,23
 chalk/cli.py,sha256=ckqqfOI-A2mT23-rnZzDMmblYj-2x1VBX8ebHlIEn9A,5873
 chalk/importer.py,sha256=m4lMn1lSYj_euDq8CS7LYTBnek9JOcjGJf9-82dJHbA,64441
 chalk/prompts.py,sha256=2H9UomLAamdfRTNUdKs9i3VTpiossuyRhntqsAXUhhg,16117
@@ -616,10 +616,10 @@ chalk/client/client.py,sha256=VfaLfbWyk-VaImPLNrNeAP4FvuJwQ5yEtJrkLST_-H4,104808
 chalk/client/client_async.py,sha256=YmuExumlDkenlGc8ROtmPjUE2m6nQUZisUDvJmwTRdk,51546
 chalk/client/client_async_impl.py,sha256=ZphhgTB49JBWHCGXe-dI0wWWKc9zPcOczy02q_gFy50,6925
 chalk/client/client_grpc.py,sha256=z1WeFWjmKKCU8dK7fdvWEOsCUlBi-v_2-HijC2_oIVs,108854
-chalk/client/client_impl.py,sha256=L0qE7Gr5MTMv71lg7-HkoPS672eERGVsfu8FObM_ZyQ,213759
+chalk/client/client_impl.py,sha256=Izb7a1sOLLI3vK8_Mi51iPZc5kmaqPxAhbqjHvzRF90,216305
 chalk/client/dataset.py,sha256=LneWwaAOHCjtj7gaJjsSeVNruj-QJ51hjRi62zrFNVE,77561
 chalk/client/exc.py,sha256=kZJ80YbSeSRDmTLTh240j_eRdJFZBa7IaDsNSRoDroU,4145
-chalk/client/models.py,sha256=TWiRHDMP_9NUqlCAGaw2RI2UlFI5ED1BCBM1rZRK72k,67482
+chalk/client/models.py,sha256=ApXv28A44dNXUULKFkG5L6e48vDscZyFGenaXItxL94,67598
 chalk/client/response.py,sha256=m8sQCOj7YVv3mZSZMIC1rIMzFMQ9rfMdBRLg5NRmOOE,53257
 chalk/client/_internal_models/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 chalk/client/_internal_models/check.py,sha256=3Xfo4Ws4rvwjeVg0-5-kejfRfRBJeqHmnRhW-WEz784,917
@@ -660,7 +660,7 @@ chalk/features/live_updates.py,sha256=8ZbiDjcLqfFruSL15_aycwzSqJ0TbKNhas06KfZLyL
 chalk/features/namespace_context.py,sha256=fL-nPohqtNiyPDS1uQTAaHLns4aivuBL2Flf50DajU4,1813
 chalk/features/primary.py,sha256=BZ8mrMmKfRNy_wnKGDJt2cdnejP_CZb6xBsD9Ljgajc,5209
 chalk/features/pseudofeatures.py,sha256=50Pe_Xi8ttYWtgNNRpgkhBxP8xoCZCYwyLb0aWUQ-PI,2147
-chalk/features/resolver.py,sha256=Et-cngxAjqbstJC5lBD3AhMwzXvUjQ6Tj8FnnW5bACU,193610
+chalk/features/resolver.py,sha256=-VU7027xob5StsHC4_UlR1AhH_uxcW_jiwznKHHIMNw,193745
 chalk/features/tag.py,sha256=LRmKRA8ANCOvmaIAk-L5j1QW2U0aah2SeASy8Uydkmk,2675
 chalk/features/underscore.py,sha256=4xnfQV3bfvVn0PNEtkT4J-k7hW4ebtH9KBe4_BvGjY4,26763
 chalk/features/underscore_features.py,sha256=PlVCoaDDffOgtiSMaxPSWCoj8IjscbkOzDLA471HsJ4,13005
@@ -827,8 +827,8 @@ chalk/utils/tracing.py,sha256=NiiM-9dbuJhSCv6R1npR1uYNKWlkqTR6Ygm0Voi2NrY,13078
 chalk/utils/weak_set_by_identity.py,sha256=VmikA_laYwFeOphCwXJIuyOIkrdlQe0bSzaXq7onoQw,953
 chalk/utils/pydanticutil/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 chalk/utils/pydanticutil/pydantic_compat.py,sha256=O575lLYJ5GvZC4HMzR9yATxf9XwjC6NrDUXbNwZidlE,3031
-chalkpy-2.98.6.dist-info/METADATA,sha256=_FtE0_mtvo_reEoEYa0_LgLU4Dn5-3vO0Y20jvu-4sU,27754
-chalkpy-2.98.6.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
-chalkpy-2.98.6.dist-info/entry_points.txt,sha256=Vg23sd8icwq-morJrljVFr-kQnMbm95rZfZj5wsZGis,42
-chalkpy-2.98.6.dist-info/top_level.txt,sha256=1Q6_19IGYfNxSw50W8tYKEJ2t5HKQ3W9Wiw4ia5yg2c,6
-chalkpy-2.98.6.dist-info/RECORD,,
+chalkpy-2.99.1.dist-info/METADATA,sha256=0h_o7NzQ0REbB61J16f_NpXpes2lu_0S-trfMzb5ZdQ,27754
+chalkpy-2.99.1.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+chalkpy-2.99.1.dist-info/entry_points.txt,sha256=Vg23sd8icwq-morJrljVFr-kQnMbm95rZfZj5wsZGis,42
+chalkpy-2.99.1.dist-info/top_level.txt,sha256=1Q6_19IGYfNxSw50W8tYKEJ2t5HKQ3W9Wiw4ia5yg2c,6
+chalkpy-2.99.1.dist-info/RECORD,,