chalkpy 2.92.3__tar.gz → 2.92.5__tar.gz

{chalkpy-2.92.3 → chalkpy-2.92.5}/PKG-INFO

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

chalkpy-2.92.5/chalk/_version.py

@@ -0,0 +1 @@
+__version__ = "2.92.5"

{chalkpy-2.92.3 → chalkpy-2.92.5}/chalk/client/client_async_impl.py

@@ -161,7 +161,7 @@ class AsyncChalkClientImpl(AsyncChalkClient):
     ) -> BulkOnlineQueryResult:
         return await self._current_loop().run_in_executor(
             self._executor,
-            partial(self._client._run_serialized_query, *args, **kwargs),
+            partial(self._client._run_serialized_query, *args, **kwargs),  # pyright: ignore[reportPrivateUsage]
         )
 
     async def offline_query(

{chalkpy-2.92.3 → chalkpy-2.92.5}/chalk/client/client_grpc.py

@@ -869,7 +869,7 @@ class ChalkGRPCClient:
 
     def online_query_bulk(
         self,
-        input: Union[Mapping[FeatureReference, Sequence[Any]], DataFrame, str],
+        input: Union[Mapping[FeatureReference, Sequence[Any]], DataFrame, None] = None,
         output: Sequence[FeatureReference] = (),
         now: Optional[Sequence[dt.datetime]] = None,
         staleness: Optional[Mapping[FeatureReference, str]] = None,
@@ -888,9 +888,17 @@ class ChalkGRPCClient:
         request_timeout: Optional[float] = None,
         headers: Mapping[str, str | bytes] | Sequence[tuple[str, str | bytes]] | None = None,
         query_context: Mapping[str, Union[str, int, float, bool, None]] | str | None = None,
+        *,
+        input_sql: str | None = None,
     ) -> BulkOnlineQueryResult:
+        if input is None and input_sql is None:
+            raise TypeError("One of `input` or `input_sql` is required")
+        if input is not None and input_sql is not None:
+            raise TypeError("`input` and `input_sql` are mutually exclusive")
+
         response, call = self._online_query_bulk_grpc_request(
             input=input,
+            input_sql=input_sql,
             output=output,
             now=now,
             staleness=staleness,
@@ -917,7 +925,8 @@ class ChalkGRPCClient:
     def _online_query_bulk_grpc_request(
         self,
         *,
-        input: Union[Mapping[FeatureReference, Sequence[Any]], DataFrame, str],
+        input: Union[Mapping[FeatureReference, Sequence[Any]], DataFrame, None] = None,
+        input_sql: str | None = None,
         output: Sequence[FeatureReference] = (),
         now: Optional[Sequence[dt.datetime]] = None,
         staleness: Optional[Mapping[FeatureReference, str]] = None,
@@ -938,8 +947,10 @@ class ChalkGRPCClient:
         query_context: Mapping[str, Union[str, int, float, bool, None]] | None = None,
     ) -> Tuple[online_query_pb2.OnlineQueryBulkResponse, grpc.Call]:
         """Returns the raw GRPC response and metadata"""
+
         request = self._make_query_bulk_request(
             input=input,
+            input_sql=input_sql,
             output=output,
             now=now or (),
             staleness=staleness or {},
@@ -1128,7 +1139,9 @@ class ChalkGRPCClient:
 
     def _make_query_bulk_request(
         self,
-        input: Mapping[FeatureReference, Sequence[Any]] | DataFrame | str,
+        *,
+        input: Mapping[FeatureReference, Sequence[Any]] | DataFrame | None = None,
+        input_sql: str | None = None,
         output: Sequence[FeatureReference],
         now: Sequence[dt.datetime],
         staleness: Mapping[FeatureReference, str],
@@ -1146,18 +1159,18 @@ class ChalkGRPCClient:
         planner_options: Mapping[str, str | int | bool],
         query_context: Mapping[str, Union[str, int, float, bool, None]] | str | None,
     ) -> online_query_pb2.OnlineQueryBulkRequest:
+        if input is None and input_sql is None:
+            raise TypeError("One of `input` or `input_sql` is required")
+        if input is not None and input_sql is not None:
+            raise TypeError("`input` and `input_sql` are mutually exclusive")
+
         inputs_feather: bytes | None
-        inputs_sql: str | None
-        if isinstance(input, str):
-            # It's a SQL query.
+        if input is None:
             inputs_feather = None
-            inputs_sql = input
         else:
-            # It's literal data.
             inputs_feather = get_features_feather_bytes(
                 input, self._INPUT_ENCODE_OPTIONS, compression=self._input_compression
             )
-            inputs_sql = None
 
         encoded_outputs = encode_outputs(output)
         outputs = encoded_outputs.string_outputs
@@ -1188,7 +1201,7 @@ class ChalkGRPCClient:
         query_context_proto = {k: value_to_proto(v) for k, v in query_context.items()} if query_context else None
         return online_query_pb2.OnlineQueryBulkRequest(
             inputs_feather=inputs_feather,
-            inputs_sql=inputs_sql,
+            inputs_sql=input_sql,
             outputs=[online_query_pb2.OutputExpr(feature_fqn=o) for o in outputs]
             + [online_query_pb2.OutputExpr(feature_expression=o) for o in encoded_outputs.feature_expressions_proto],
             now=now_proto,

{chalkpy-2.92.3 → chalkpy-2.92.5}/chalk/client/client_impl.py

@@ -2357,9 +2357,9 @@ https://docs.chalk.ai/cli/apply
             override_target_image_tag=override_target_image_tag,
             num_shards=num_shards,
             num_workers=num_workers,
-            feature_for_lower_upper_bound=str(feature_for_lower_upper_bound)
-            if feature_for_lower_upper_bound is not None
-            else None,
+            feature_for_lower_upper_bound=(
+                str(feature_for_lower_upper_bound) if feature_for_lower_upper_bound is not None else None
+            ),
             completion_deadline=completion_deadline,
             max_retries=max_retries,
             optional_output_expressions=optional_output_expressions,
@@ -4419,7 +4419,7 @@ https://docs.chalk.ai/cli/apply
         """
         try:
             import pyarrow as pa
-            import pyarrow.feather
+            import pyarrow.feather as feather
         except ImportError:
             raise missing_dependency_exception("chalkpy[runtime]")
 
@@ -4452,6 +4452,7 @@ https://docs.chalk.ai/cli/apply
 
         # Create FeatherRequestHeader
         from chalk.client.models import OnlineQueryContext as OQC
+
         header_dict = {
             "outputs": outputs_encoded,
             "expression_outputs": [],
@@ -4469,7 +4470,7 @@ https://docs.chalk.ai/cli/apply
 
         # Serialize the input table to feather format
         feather_buffer = BytesIO()
-        pa.feather.write_feather(input_table, feather_buffer)
+        feather.write_feather(input_table, feather_buffer)
         feather_bytes = feather_buffer.getvalue()
 
         # Build the request body:
@@ -4495,13 +4496,13 @@ https://docs.chalk.ai/cli/apply
             json=None,
             data=request_body.getvalue(),
             environment_override=None,
-            preview_deployment_id=...,
-            branch=...,
+            preview_deployment_id=None,
+            branch=None,
             metadata_request=False,
             extra_headers=headers,
         )
 
-        if not isinstance(response, requests.Response):
+        if not isinstance(response, requests.Response):  # pyright: ignore[reportUnnecessaryIsInstance]
             raise TypeError("Expected requests.Response")
 
         if response.status_code != 200:
@@ -4513,7 +4514,7 @@ https://docs.chalk.ai/cli/apply
         # Convert feather bytes back to a dataframe
         scalars_df = None
         if result.scalar_data:
-            scalars_table = pa.feather.read_table(BytesIO(result.scalar_data))
+            scalars_table = feather.read_table(BytesIO(result.scalar_data))
             scalars_df = pa_table_to_pl_df(scalars_table)
 
         # Parse errors from JSON strings back to ChalkError objects

{chalkpy-2.92.3 → chalkpy-2.92.5}/chalk/sql/__init__.py

@@ -1437,6 +1437,8 @@ def DatabricksSource(
     access_token: str = ...,
     db: str = ...,
     port: str = ...,
+    client_id: str = ...,
+    client_secret: str = ...,
     engine_args: Optional[Dict[str, Any]] = ...,
 ) -> BaseSQLSourceProtocol:
     """You can also configure the integration directly using environment
@@ -1457,6 +1459,10 @@ def DatabricksSource(
         Database to use.
     port
         Port number to use.
+    client_id
+        OAuth service principal client ID (alternative to access_token).
+    client_secret
+        OAuth service principal client secret (alternative to access_token).
     engine_args
         Additional arguments to use when constructing the SQLAlchemy engine.
 
@@ -1475,6 +1481,14 @@ def DatabricksSource(
     ...     db=os.getenv("DATABRICKS_DATABASE"),
     ...     port=os.getenv("DATABRICKS_PORT"),
     ... )
+    >>> databricks_with_oauth = DatabricksSource(
+    ...     host=os.getenv("DATABRICKS_HOST"),
+    ...     http_path=os.getenv("DATABRICKS_HTTP_PATH"),
+    ...     client_id=os.getenv("DATABRICKS_CLIENT_ID"),
+    ...     client_secret=os.getenv("DATABRICKS_CLIENT_SECRET"),
+    ...     db=os.getenv("DATABRICKS_DATABASE"),
+    ...     port=os.getenv("DATABRICKS_PORT"),
+    ... )
     """
     ...
 
@@ -1487,6 +1501,8 @@ def DatabricksSource(
     access_token: Optional[str] = None,
     db: Optional[str] = None,
     port: Optional[Union[str, int]] = None,
+    client_id: Optional[str] = None,
+    client_secret: Optional[str] = None,
     engine_args: Optional[Dict[str, Any]] = None,
 ) -> BaseSQLSourceProtocol:
     """Create a Databricks data source. SQL-based data sources
@@ -1505,6 +1521,8 @@ def DatabricksSource(
         db=db,
         port=port,
         name=name,
+        client_id=client_id,
+        client_secret=client_secret,
         engine_args=engine_args,
     )
 

{chalkpy-2.92.3 → chalkpy-2.92.5}/chalk/sql/_internal/integrations/databricks.py

@@ -33,6 +33,8 @@ _DATABRICKS_HTTP_PATH_NAME = "DATABRICKS_HTTP_PATH"
 _DATABRICKS_TOKEN_NAME = "DATABRICKS_TOKEN"
 _DATABRICKS_DATABASE_NAME = "DATABRICKS_DATABASE"
 _DATABRICKS_PORT_NAME = "DATABRICKS_PORT"
+_DATABRICKS_CLIENT_ID_NAME = "DATABRICKS_CLIENT_ID"
+_DATABRICKS_CLIENT_SECRET_NAME = "DATABRICKS_CLIENT_SECRET"
 
 
 class DatabricksSourceImpl(BaseSQLSource):
@@ -46,6 +48,8 @@ class DatabricksSourceImpl(BaseSQLSource):
         db: Optional[str] = None,
         port: Optional[Union[int, str]] = None,
         name: Optional[str] = None,
+        client_id: Optional[str] = None,
+        client_secret: Optional[str] = None,
         engine_args: Optional[Dict[str, Any]] = None,
         executor: Optional[concurrent.futures.ThreadPoolExecutor] = None,
         integration_variable_override: Optional[Mapping[str, str]] = None,
@@ -74,21 +78,53 @@ class DatabricksSourceImpl(BaseSQLSource):
                 name=_DATABRICKS_PORT_NAME, integration_name=name, parser=int, override=integration_variable_override
             )
         )
+        self.client_id = client_id or load_integration_variable(
+            name=_DATABRICKS_CLIENT_ID_NAME, integration_name=name, override=integration_variable_override
+        )
+        self.client_secret = client_secret or load_integration_variable(
+            name=_DATABRICKS_CLIENT_SECRET_NAME, integration_name=name, override=integration_variable_override
+        )
         self.executor = executor or DEFAULT_IO_EXECUTOR
 
+        has_token = self.access_token is not None
+        has_oauth = self.client_id is not None and self.client_secret is not None
+
+        if has_token and has_oauth:
+            chalk_logger.warning(
+                "Both OAuth credentials and a personal access token were provided. Using OAuth authentication."
+            )
+            self.access_token = None
+
+        self._credentials_provider = None
+        if has_oauth:
+            try:
+                from databricks.sdk.core import Config, oauth_service_principal
+            except ImportError:
+                raise missing_dependency_exception("chalkpy[databricks]")
+
+            def credentials_provider():
+                config = Config(host=self.host, client_id=self.client_id, client_secret=self.client_secret)
+                return oauth_service_principal(config)
+
+            self._credentials_provider = credentials_provider
+
         if engine_args is None:
             engine_args = {}
+
+        connect_args: dict[str, Any] = {
+            "keepalives": 1,
+            "keepalives_idle": 30,
+            "keepalives_interval": 10,
+            "keepalives_count": 5,
+        }
+
+        if self._credentials_provider:
+            connect_args["credentials_provider"] = self._credentials_provider
+
         engine_args.setdefault("pool_size", 20)
         engine_args.setdefault("max_overflow", 60)
-        engine_args.setdefault(
-            "connect_args",
-            {
-                "keepalives": 1,
-                "keepalives_idle": 30,
-                "keepalives_interval": 10,
-                "keepalives_count": 5,
-            },
-        )
+        engine_args.setdefault("connect_args", connect_args)
+
         BaseSQLSource.__init__(self, name=name, engine_args=engine_args, async_engine_args={})
 
     def supports_inefficient_fallback(self) -> bool:
@@ -97,6 +133,12 @@ class DatabricksSourceImpl(BaseSQLSource):
     def get_sqlglot_dialect(self) -> str | None:
         return "databricks"
 
+    def _get_client_auth(self) -> Dict[str, str | Callable | None]:
+        if self._credentials_provider:
+            return {"credentials_provider": self._credentials_provider}
+        else:
+            return {"access_token": self.access_token}
+
     @contextlib.contextmanager
     def _create_temp_table(
         self,
@@ -168,10 +210,7 @@ class DatabricksSourceImpl(BaseSQLSource):
 
                     # Connect using databricks-sql-python for efficient Arrow fetching
                     with sql.connect(
-                        server_hostname=self.host,
-                        http_path=self.http_path,
-                        access_token=self.access_token,
-                        catalog=self.db,
+                        server_hostname=self.host, http_path=self.http_path, catalog=self.db, **self._get_client_auth()
                     ) as databricks_conn:
                         chalk_logger.info("Established connection with Databricks using databricks-sql-python")
 
@@ -273,10 +312,7 @@ class DatabricksSourceImpl(BaseSQLSource):
 
             # Connect using databricks-sql-python for efficient Arrow fetching
             with sql.connect(
-                server_hostname=self.host,
-                http_path=self.http_path,
-                access_token=self.access_token,
-                catalog=self.db,
+                server_hostname=self.host, http_path=self.http_path, catalog=self.db, **self._get_client_auth()
             ) as databricks_cnx:
                 with databricks_cnx.cursor() as cursor:
                     formatted_op, positional_params, named_params = self.compile_query(finalized_query)
@@ -353,6 +389,8 @@ class DatabricksSourceImpl(BaseSQLSource):
                 create_integration_variable(_DATABRICKS_TOKEN_NAME, self.name, self.access_token),
                 create_integration_variable(_DATABRICKS_DATABASE_NAME, self.name, self.db),
                 create_integration_variable(_DATABRICKS_PORT_NAME, self.name, self.port),
+                create_integration_variable(_DATABRICKS_CLIENT_ID_NAME, self.name, self.client_id),
+                create_integration_variable(_DATABRICKS_CLIENT_SECRET_NAME, self.name, self.client_secret),
             ]
             if v is not None
         }

chalkpy-2.92.3/chalk/_version.py

@@ -1 +0,0 @@
-__version__ = "2.92.3"