chalkpy 2.98.1__py3-none-any.whl → 2.98.2__py3-none-any.whl

chalk/_version.py

@@ -1 +1 @@
-__version__ = "2.98.1"
+__version__ = "2.98.2"

chalk/client/client_impl.py

@@ -1033,10 +1033,6 @@ 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",
@@ -1047,7 +1043,7 @@ class ChalkAPIClientImpl(ChalkClient):
         if additional_headers:
             self._default_headers.update(additional_headers)
 
-        self._primary_environment: EnvironmentId | None = environment
+        self._primary_environment = environment
 
         self.__class__.latest_client = self
         if notebook.is_notebook():
@@ -1120,16 +1116,11 @@ class ChalkAPIClientImpl(ChalkClient):
             creds = ExchangeCredentialsResponse(**response_json)
         except ValidationError:
             raise HTTPError(response=resp)
-        self._access_token = creds.access_token
         self._default_headers["Authorization"] = f"Bearer {creds.access_token}"
         # 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
         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,
@@ -1392,48 +1383,6 @@ 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 | None:
-        """
-        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:
-            return self._api_server
-        else:
-            return self._get_engine_host(environment_override)
-
     def _request(
         self,
         method: str,
@@ -1450,20 +1399,25 @@ 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.
         allow_credential_exchange: bool = True
-        if self._access_token is None:
+
+        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
             self._exchange_credentials()
             allow_credential_exchange = False
 
-        host = self._get_request_host(metadata_request, environment_override)
+            # After exchanging credentials, the api server is never none
+            assert self._api_server is not None
 
-        # After exchanging credentials, the api server is never null, and the request host
-        # helper function falls back to the api server if it cannot determine an engine url
-        assert host is not None, "After credential exchange, the request host must be non-null"
+            if metadata_request or self._query_server is None:
+                host = self._api_server
+            else:
+                host = self._query_server
 
         r = self._do_request_inner(
             method=method,
@@ -1484,11 +1438,14 @@ 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()
-            host = self._get_request_host(metadata_request, environment_override)
 
-            # After exchanging credentials, the api server is never null, and the request host
-            # helper function falls back to the api server if it cannot determine an engine url
-            assert host is not None, "After credential exchange, the request host must be non-null"
+            # 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
 
             r = self._do_request_inner(
                 method=method,
@@ -4030,8 +3987,9 @@ https://docs.chalk.ai/cli/apply
             environment_override=context.environment,
             preview_deployment_id=preview_deployment_id,
             branch=branch,
-            # all offline query requests should be routed through api server
-            metadata_request=True,
+            # 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,
         )
         return response
 
@@ -4107,6 +4065,19 @@ 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,8 +454,6 @@ 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):

{chalkpy-2.98.1.dist-info → chalkpy-2.98.2.dist-info}/METADATA

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

{chalkpy-2.98.1.dist-info → chalkpy-2.98.2.dist-info}/RECORD

@@ -1,5 +1,5 @@
 chalk/__init__.py,sha256=vKsx9-cl5kImlVWGHVRYO6bweBm79NAzGs3l36u71wM,2657
-chalk/_version.py,sha256=UJbPCGwo0rdAj6Ro_XeyrGHCc_AZqBH9dWB9R1_fZVQ,23
+chalk/_version.py,sha256=7Cc2OCS4iZEpzdq7iHwkSLQA8S6psTq-SKpiAAtUbG0,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=DcRJltVqWy_B3Qvbnjz3qOECT9TZ7KD0eovds2CNRIw,104746
 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=0z6f9LK4t39VNazfJ9c140Szq3R1TcZlPe_uQh20wFw,216225
+chalk/client/client_impl.py,sha256=m2-dovRppOchuEwt2LlKJAv0LBW5obs9ztuJeDhuqY0,213660
 chalk/client/dataset.py,sha256=LneWwaAOHCjtj7gaJjsSeVNruj-QJ51hjRi62zrFNVE,77561
 chalk/client/exc.py,sha256=kZJ80YbSeSRDmTLTh240j_eRdJFZBa7IaDsNSRoDroU,4145
-chalk/client/models.py,sha256=ApXv28A44dNXUULKFkG5L6e48vDscZyFGenaXItxL94,67598
+chalk/client/models.py,sha256=TWiRHDMP_9NUqlCAGaw2RI2UlFI5ED1BCBM1rZRK72k,67482
 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
@@ -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.1.dist-info/METADATA,sha256=kcf9CjMFEkKTs35bwUw0UfJu-8foHPLqXHwlkiA7heU,27754
-chalkpy-2.98.1.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
-chalkpy-2.98.1.dist-info/entry_points.txt,sha256=Vg23sd8icwq-morJrljVFr-kQnMbm95rZfZj5wsZGis,42
-chalkpy-2.98.1.dist-info/top_level.txt,sha256=1Q6_19IGYfNxSw50W8tYKEJ2t5HKQ3W9Wiw4ia5yg2c,6
-chalkpy-2.98.1.dist-info/RECORD,,
+chalkpy-2.98.2.dist-info/METADATA,sha256=iROjKSQYMlH8Dfat6xKGCMx7ymmRV7SvKhEK_yGolb8,27754
+chalkpy-2.98.2.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+chalkpy-2.98.2.dist-info/entry_points.txt,sha256=Vg23sd8icwq-morJrljVFr-kQnMbm95rZfZj5wsZGis,42
+chalkpy-2.98.2.dist-info/top_level.txt,sha256=1Q6_19IGYfNxSw50W8tYKEJ2t5HKQ3W9Wiw4ia5yg2c,6
+chalkpy-2.98.2.dist-info/RECORD,,