chalkpy 2.95.2__tar.gz → 2.95.3__tar.gz

{chalkpy-2.95.2 → chalkpy-2.95.3}/PKG-INFO

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

chalkpy-2.95.3/chalk/_version.py

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

{chalkpy-2.95.2 → chalkpy-2.95.3}/chalk/sql/_internal/integrations/mssql.py

@@ -2,15 +2,19 @@ from __future__ import annotations
 
 import os
 import struct
-from typing import TYPE_CHECKING, Any, Dict, Mapping, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, Iterable, Mapping, Optional, Union
 
 from chalk.integrations.named import create_integration_variable, load_integration_variable
+from chalk.sql._internal.query_execution_parameters import QueryExecutionParameters
 from chalk.sql._internal.sql_source import BaseSQLSource, SQLSourceKind, TableIngestMixIn
+from chalk.sql.finalized_query import FinalizedChalkQuery
 from chalk.sql.protocols import SQLSourceWithTableIngestProtocol
+from chalk.utils.environment_parsing import env_var_bool
 from chalk.utils.missing_dependency import missing_dependency_exception
 
 if TYPE_CHECKING:
-    from sqlalchemy.engine import URL
+    import pyarrow as pa
+    from sqlalchemy.engine import URL, Connection
 
 _MSSQL_HOST_NAME = "MSSQL_HOST"
 _MSSQL_TCP_PORT_NAME = "MSSQL_TCP_PORT"
@@ -42,8 +46,8 @@ class MSSQLSourceImpl(BaseSQLSource, TableIngestMixIn, SQLSourceWithTableIngestP
     ):
         try:
             import pyodbc
-        except ImportError:
-            raise missing_dependency_exception("chalkpy[mssql]")
+        except ImportError as e:
+            raise missing_dependency_exception("chalkpy[mssql]", original_error=e)
         del pyodbc
 
         self.name = name
@@ -142,9 +146,17 @@ class MSSQLSourceImpl(BaseSQLSource, TableIngestMixIn, SQLSourceWithTableIngestP
     def local_engine_url(self) -> "URL":
         from sqlalchemy.engine.url import URL
 
+        trust_server_cert = env_var_bool("CHALK_MSSQL_TRUST_SERVER_CERTIFICATE", default=False)
+
         if self.client_id and self.client_secret and self.tenant_id:
             # Service Principal authentication
             # Use pyodbc driver for Azure AD support
+            query_params = {
+                "driver": "ODBC Driver 18 for SQL Server",
+                "Authentication": "ActiveDirectoryServicePrincipal",
+            }
+            if trust_server_cert:
+                query_params["TrustServerCertificate"] = "yes"
             return URL.create(
                 drivername="mssql+pyodbc",
                 username=self.client_id,
@@ -152,13 +164,13 @@ class MSSQLSourceImpl(BaseSQLSource, TableIngestMixIn, SQLSourceWithTableIngestP
                 host=self.host,
                 port=self.port,
                 database=self.db,
-                query={
-                    "driver": "ODBC Driver 18 for SQL Server",
-                    "Authentication": "ActiveDirectoryServicePrincipal",
-                },
+                query=query_params,
             )
         elif self.user and self.password:
             # SQL authentication
+            query_params = {"driver": "ODBC Driver 18 for SQL Server"}
+            if trust_server_cert:
+                query_params["TrustServerCertificate"] = "yes"
             return URL.create(
                 drivername="mssql+pyodbc",
                 username=self.user,
@@ -166,13 +178,15 @@ class MSSQLSourceImpl(BaseSQLSource, TableIngestMixIn, SQLSourceWithTableIngestP
                 host=self.host,
                 port=self.port,
                 database=self.db,
-                query={"driver": "ODBC Driver 18 for SQL Server"},
+                query=query_params,
             )
         else:
             # Managed Identity: token injected via event listener
             connection_string = (
                 f"DRIVER={{ODBC Driver 18 for SQL Server}};SERVER={self.host},{self.port};DATABASE={self.db}"
             )
+            if trust_server_cert:
+                connection_string += ";TrustServerCertificate=yes"
             return URL.create(
                 drivername="mssql+pyodbc",
                 query={"odbc_connect": connection_string},
@@ -193,3 +207,75 @@ class MSSQLSourceImpl(BaseSQLSource, TableIngestMixIn, SQLSourceWithTableIngestP
             ]
             if v is not None
         }
+
+    def execute_query_efficient_raw(
+        self,
+        finalized_query: FinalizedChalkQuery,
+        expected_output_schema: "pa.Schema",
+        connection: Optional["Connection"],
+        query_execution_parameters: QueryExecutionParameters,
+    ) -> Iterable["pa.RecordBatch"]:
+        """Execute query efficiently for MSSQL and return raw PyArrow RecordBatches."""
+        import contextlib
+
+        import pyarrow as pa
+        import pyarrow.compute as pc
+
+        # Get the compiled query
+        _, _, _ = self.compile_query(finalized_query)
+
+        # Use existing connection or create new one
+        with (self.get_engine().connect() if connection is None else contextlib.nullcontext(connection)) as cnx:
+            with cnx.begin():
+                # Handle temp tables
+                with contextlib.ExitStack() as exit_stack:
+                    for (
+                        _,
+                        temp_value,
+                        create_temp_table,
+                        temp_table,
+                        drop_temp_table,
+                    ) in finalized_query.temp_tables.values():
+                        exit_stack.enter_context(
+                            self._create_temp_table(create_temp_table, temp_table, drop_temp_table, cnx, temp_value)
+                        )
+
+                    # Execute query
+                    result = cnx.execute(finalized_query.query, finalized_query.params)
+
+                    # Convert result to PyArrow
+                    rows = result.fetchall()
+                    column_names = result.keys()
+
+                    if not rows:
+                        # Return empty batch with expected schema
+                        arrays = [pa.nulls(0, field.type) for field in expected_output_schema]
+                        batch = pa.RecordBatch.from_arrays(arrays, schema=expected_output_schema)
+                        if query_execution_parameters.yield_empty_batches:
+                            yield batch
+                        return
+
+                    # Convert rows to column arrays
+                    data: dict[str, list[Any]] = {}
+                    for i, col_name in enumerate(column_names):
+                        col_data = [row[i] for row in rows]
+                        data[col_name] = col_data
+
+                    # Create PyArrow table
+                    table = pa.table(data)
+
+                    # Map columns to expected schema
+                    arrays: list[pa.Array] = []
+                    for field in expected_output_schema:
+                        if field.name in table.column_names:
+                            col = table.column(field.name)
+                            # Cast to expected type if needed
+                            if col.type != field.type:
+                                col = pc.cast(col, field.type)
+                            arrays.append(col)
+                        else:
+                            # Column not found, create null array
+                            arrays.append(pa.nulls(len(table), field.type))
+
+                    batch = pa.RecordBatch.from_arrays(arrays, schema=expected_output_schema)
+                    yield batch

chalkpy-2.95.3/chalk/utils/missing_dependency.py

@@ -0,0 +1,12 @@
+import chalk
+
+
+class MissingDependencyException(ImportError):
+    ...
+
+
+def missing_dependency_exception(name: str, original_error: Exception | None = None):
+    msg = f"Missing pip dependency '{name}' for chalkpy=={chalk.__version__}. Please add this to your requirements.txt file and pip install."
+    if original_error:
+        msg += f"\n\n{original_error}"
+    return MissingDependencyException(msg)

chalkpy-2.95.2/chalk/_version.py

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

chalkpy-2.95.2/chalk/utils/missing_dependency.py

@@ -1,11 +0,0 @@
-import chalk
-
-
-class MissingDependencyException(ImportError):
-    ...
-
-
-def missing_dependency_exception(name: str):
-    return MissingDependencyException(
-        f"Missing pip dependency '{name}' for chalkpy=={chalk.__version__}. Please add this to your requirements.txt file and pip install."
-    )