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

chalk/_version.py

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

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):

chalk/operators/_utils.py

@@ -2,11 +2,12 @@ from __future__ import annotations
 
 import dataclasses
 import traceback
-from typing import Callable, Optional, Sequence, Union
+from typing import Callable, Optional, Protocol, Sequence, TypeVar, Union
 
 import pyarrow
 
 from chalk import DataFrame, Features, StaticOperator
+from chalk._gen.chalk.dataframe.v1 import dataframe_pb2
 from chalk._gen.chalk.expression.v1 import expression_pb2 as expr_pb
 from chalk.client import ChalkError, ChalkException, ErrorCode, ErrorCodeCategory
 from chalk.df.LazyFramePlaceholder import LazyFramePlaceholder
@@ -40,6 +41,14 @@ class _GetStaticOperatorError(Exception):
         )
 
 
+ProtoT = TypeVar("ProtoT", covariant=True)
+
+
+class _ToProto(Protocol[ProtoT]):
+    def _to_proto(self) -> ProtoT:
+        ...
+
+
 @dataclasses.dataclass
 class DfPlaceholder:
     schema_dict: dict[str, pyarrow.DataType]
@@ -80,7 +89,7 @@ def static_resolver_to_operator(
     fn: Callable,
     inputs: Sequence[Union[Feature, type[DataFrame]]],
     output: Optional[type[Features]],
-) -> StaticOperator | DfPlaceholder | ChalkDataFrame | LazyFramePlaceholder:
+) -> StaticOperator | DfPlaceholder | ChalkDataFrame | _ToProto[dataframe_pb2.DataFramePlan]:
     if output is None:
         raise _GetStaticOperatorError(
             resolver_fqn=fqn,

chalk/parsed/to_proto.py

@@ -16,12 +16,13 @@ from typing_extensions import assert_never
 
 from chalk import DataFrame
 from chalk._gen.chalk.arrow.v1 import arrow_pb2 as arrow_pb
+from chalk._gen.chalk.dataframe.v1.dataframe_pb2 import DataFramePlan
 from chalk._gen.chalk.expression.v1 import expression_pb2 as expr_pb
+from chalk._gen.chalk.expression.v1.expression_pb2 import LogicalExprNode
 from chalk._gen.chalk.graph.v1 import graph_pb2 as pb
 from chalk._gen.chalk.graph.v2 import sources_pb2 as sources_pb
 from chalk._gen.chalk.lsp.v1.lsp_pb2 import Location, Position, Range
 from chalk._validation.feature_validation import FeatureValidation
-from chalk.df.LazyFramePlaceholder import LazyFramePlaceholder
 from chalk.features import (
     CacheStrategy,
     Feature,
@@ -1154,10 +1155,11 @@ class ToProtoConverter:
         static_operation_dataframe = None
         if r.static:
             static_operator = static_resolver_to_operator(fqn=r.fqn, fn=r.fn, inputs=r.inputs, output=r.output)
-            if isinstance(static_operator, LazyFramePlaceholder):
-                static_operation_dataframe = static_operator._to_proto()  # pyright: ignore[reportPrivateUsage]
-            else:
-                static_operation = static_operator._to_proto()  # pyright: ignore[reportPrivateUsage]
+            to_proto = static_operator._to_proto()  # pyright: ignore[reportPrivateUsage]
+            if isinstance(to_proto, DataFramePlan):
+                static_operation_dataframe = to_proto
+            elif isinstance(to_proto, LogicalExprNode):  # pyright: ignore[reportUnnecessaryIsInstance]
+                static_operation = to_proto
 
         function_reference_proto = ToProtoConverter.create_function_reference(
             r.fn,

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

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

{chalkpy-2.98.1.dist-info → chalkpy-2.98.3.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=Bz_5zhHPwDjTPM5Kpa6MCRtG34E0bkBIbT_od0j6sq8,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
@@ -718,7 +718,7 @@ chalk/operators/_literal.py,sha256=T8Euf2HRUVlFuct3svX4Z8EVf80r0cWDBZGr0MUgS9Q,5
 chalk/operators/_parquet_scan.py,sha256=3pSd1U-9My145GO1Nr1P_T2SfdhPZW2PLeMvKg6l4Uo,545
 chalk/operators/_rename.py,sha256=YgpMnfYsln88gfp4AwhiI8JX58hadDvEcFYrMCcVl1o,741
 chalk/operators/_select.py,sha256=0vXxnlSlp5DVnO7WcjJJuDb8jnwRtyVyVbdP_U4gtYs,1029
-chalk/operators/_utils.py,sha256=7NRUNqFoJoFCTEbAJs4H98KRbDJshwN0RoyjPfmyzN8,5050
+chalk/operators/_utils.py,sha256=WZRYAlK3M_w5k_4tyqg0AzR1cntqfaR7upNmia-g6z4,5270
 chalk/parsed/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 chalk/parsed/_graph_validation.py,sha256=XHVPw4cf-nkPdHrqGXevsFuOrMtWtibUTRkTuizEc4g,23771
 chalk/parsed/_metadata.py,sha256=TCjKr20Fs4Yguo97yKjATurwOdV-DoFQeiGxbX_s1hs,6749
@@ -727,7 +727,7 @@ chalk/parsed/branch_state_rich.py,sha256=X9st1vRMRMOfnEy3hwitsAyNcTzZlubXJt9jI1M
 chalk/parsed/duplicate_input_gql.py,sha256=IbVRKDCS-M7f9k127LpeVjRlXvzp1tuIek-bOqTULC4,20196
 chalk/parsed/expressions.py,sha256=A8U54v9jacLX4UVtLqSUYUjL-lOnI-kWvG2hHZvfYE0,554
 chalk/parsed/json_conversions.py,sha256=GRBg2KvwahP1tXwr7n--s7ueULN6prEDPO2I9B9MWzk,30459
-chalk/parsed/to_proto.py,sha256=ZSwEYkuwr3UhMVabt6FM02fBN_oBy_mSMjDYS_crJw0,75530
+chalk/parsed/to_proto.py,sha256=JsMnLlkXL1KH4LxZJStkwBHOB2geG9hnO3ku4LTl3vU,75657
 chalk/parsed/user_types_to_json.py,sha256=ZJWdYFqyhr5InvItQybtHadXQjW3vjHrv8hjMGtL3Bc,13318
 chalk/parsed/validation_from_registries.py,sha256=nfiAj1tvWRu0RrkhkGtElhAsL8V7ayEKUisrKQF-wYc,7900
 chalk/parsed/_proto/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -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.3.dist-info/METADATA,sha256=rhMVAP_p_UvatDRbPlTqI6QNHtIg3fNOioqUZ3HCwJ0,27754
+chalkpy-2.98.3.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+chalkpy-2.98.3.dist-info/entry_points.txt,sha256=Vg23sd8icwq-morJrljVFr-kQnMbm95rZfZj5wsZGis,42
+chalkpy-2.98.3.dist-info/top_level.txt,sha256=1Q6_19IGYfNxSw50W8tYKEJ2t5HKQ3W9Wiw4ia5yg2c,6
+chalkpy-2.98.3.dist-info/RECORD,,