chalkpy 2.89.22__py3-none-any.whl → 2.95.3__py3-none-any.whl

chalk/sql/__init__.py

@@ -10,6 +10,7 @@ from chalk.sql._internal.integrations.clickhouse import ClickhouseSourceImpl
 from chalk.sql._internal.integrations.cloudsql import CloudSQLSourceImpl
 from chalk.sql._internal.integrations.databricks import DatabricksSourceImpl
 from chalk.sql._internal.integrations.dynamodb import DynamoDBSourceImpl
+from chalk.sql._internal.integrations.mssql import MSSQLSourceImpl
 from chalk.sql._internal.integrations.mysql import MySQLSourceImpl
 from chalk.sql._internal.integrations.postgres import PostgreSQLSourceImpl
 from chalk.sql._internal.integrations.redshift import RedshiftSourceImpl
@@ -464,6 +465,196 @@ def MySQLSource(
     )
 
 
+@overload
+def MSSQLSource() -> SQLSourceWithTableIngestProtocol:
+    """If you have only one MSSQL connection that you'd like
+    to add to Chalk, you do not need to specify any arguments
+    to construct the source in your code.
+
+    Returns
+    -------
+    SQLSourceWithTableIngestProtocol
+        The SQL source for use in Chalk resolvers.
+
+    Examples
+    --------
+    >>> mssql = MSSQLSource()
+    """
+    ...
+
+
+@overload
+def MSSQLSource(
+    *,
+    name: str,
+    engine_args: Optional[Dict[str, Any]] = ...,
+    async_engine_args: Optional[Dict[str, Any]] = ...,
+) -> SQLSourceWithTableIngestProtocol:
+    """If you have only one MSSQL integration, there's no need to provide
+    a distinguishing name.
+
+    But what happens when you have two data sources of the same kind?
+    When you create a new data source from your dashboard,
+    you have an option to provide a name for the integration.
+    You can then reference this name in the code directly.
+
+    Parameters
+    ----------
+    name
+        Name of the integration, as configured in your dashboard.
+    engine_args
+        Additional arguments to use when constructing the SQLAlchemy engine. These arguments will be
+        merged with any default arguments from the named integration.
+    async_engine_args
+        Additional arguments to use when constructing an async SQLAlchemy engine.
+
+    Returns
+    -------
+    SQLSourceWithTableIngestProtocol
+        The SQL source for use in Chalk resolvers.
+
+    Examples
+    --------
+    >>> source = MSSQLSource(name="RISK")
+    """
+    ...
+
+
+@overload
+def MSSQLSource(
+    *,
+    name: str | None = ...,
+    host: str,
+    port: Union[int, str] = ...,
+    db: str = ...,
+    user: str = ...,
+    password: str = ...,
+    client_id: str = ...,
+    client_secret: str = ...,
+    tenant_id: str = ...,
+    engine_args: Optional[Dict[str, Any]] = ...,
+    async_engine_args: Optional[Dict[str, Any]] = ...,
+) -> SQLSourceWithTableIngestProtocol:
+    """
+    You can also configure the integration directly using environment
+    variables on your local machine or from those added through the
+    generic environment variable support (https://docs.chalk.ai/docs/env-vars).
+
+    Authentication Methods:
+    - SQL Authentication: Provide `user` and `password`
+    - Azure AD Managed Identity: Leave `user`, `password`, `client_id`, `client_secret`, and `tenant_id` empty
+    - Azure AD Service Principal: Provide `client_id`, `client_secret`, and `tenant_id`
+
+    Parameters
+    ----------
+    name
+        Name of the integration. Not required unless if this SQL Source is used within SQL File Resolvers.
+    host
+        Name of host to connect to.
+    port
+        The port number to connect to at the server host.
+    db
+        The database name.
+    user
+        MSSQL username to connect as (for SQL authentication).
+    password
+        The password to be used for SQL authentication.
+    client_id
+        Azure AD Client ID (for Service Principal authentication).
+    client_secret
+        Azure AD Client Secret (for Service Principal authentication).
+    tenant_id
+        Azure AD Tenant ID (for Service Principal authentication).
+    engine_args
+        Additional arguments to use when constructing the SQLAlchemy engine.
+    async_engine_args
+        Additional arguments to use when constructing an async SQLAlchemy engine.
+
+    Returns
+    -------
+    SQLSourceWithTableIngestProtocol
+        The SQL source for use in Chalk resolvers.
+
+    Examples
+    --------
+    SQL Authentication:
+    >>> import os
+    >>> mssql = MSSQLSource(
+    ...     host=os.getenv("MSSQL_HOST"),
+    ...     port=os.getenv("MSSQL_TCP_PORT"),
+    ...     db=os.getenv("MSSQL_DATABASE"),
+    ...     user=os.getenv("MSSQL_USER"),
+    ...     password=os.getenv("MSSQL_PWD"),
+    ... )
+
+    Managed Identity (running in Azure):
+    >>> mssql = MSSQLSource(
+    ...     host=os.getenv("MSSQL_HOST"),
+    ...     port=os.getenv("MSSQL_TCP_PORT"),
+    ...     db=os.getenv("MSSQL_DATABASE"),
+    ... )
+
+    Service Principal:
+    >>> mssql = MSSQLSource(
+    ...     host=os.getenv("MSSQL_HOST"),
+    ...     port=os.getenv("MSSQL_TCP_PORT"),
+    ...     db=os.getenv("MSSQL_DATABASE"),
+    ...     client_id=os.getenv("MSSQL_CLIENT_ID"),
+    ...     client_secret=os.getenv("MSSQL_CLIENT_SECRET"),
+    ...     tenant_id=os.getenv("MSSQL_TENANT_ID"),
+    ... )
+
+    >>> from chalk.features import online
+    >>> @online
+    ... def resolver_fn() -> User.name:
+    ...     return mssql.query_string("select name from users where id = 4").one()
+    """
+    ...
+
+
+def MSSQLSource(
+    *,
+    host: Optional[str] = None,
+    port: Optional[Union[int, str]] = None,
+    db: Optional[str] = None,
+    user: Optional[str] = None,
+    password: Optional[str] = None,
+    client_id: Optional[str] = None,
+    client_secret: Optional[str] = None,
+    tenant_id: Optional[str] = None,
+    name: Optional[str] = None,
+    engine_args: Optional[Dict[str, Any]] = None,
+    async_engine_args: Optional[Dict[str, Any]] = None,
+) -> SQLSourceWithTableIngestProtocol:
+    """Create a MSSQL data source. SQL-based data sources
+    created without arguments assume a configuration in your
+    Chalk Dashboard. Those created with the `name=` keyword
+    argument will use the configuration for the integration
+    with the given name. And finally, those created with
+    explicit arguments will use those arguments to configure
+    the data source. See the overloaded signatures for more
+    details.
+
+    Supports three authentication methods:
+    - SQL Authentication: user + password
+    - Azure AD Managed Identity: no credentials (automatic in Azure)
+    - Azure AD Service Principal: client_id + client_secret + tenant_id
+    """
+    return MSSQLSourceImpl(
+        host=host,
+        port=port,
+        db=db,
+        user=user,
+        password=password,
+        client_id=client_id,
+        client_secret=client_secret,
+        tenant_id=tenant_id,
+        name=name,
+        engine_args=engine_args,
+        async_engine_args=async_engine_args,
+    )
+
+
 def SQLiteInMemorySource(
     name: Optional[str] = None,
     engine_args: Optional[Dict[str, Any]] = None,
@@ -848,6 +1039,8 @@ def BigQuerySource(
     credentials_base64: Optional[str] = ...,
     credentials_path: Optional[str] = ...,
     engine_args: Optional[Dict[str, Any]] = ...,
+    temp_project: Optional[str] = ...,
+    temp_dataset: Optional[str] = ...,
 ) -> BaseSQLSourceProtocol:
     """You can also configure the integration directly using environment
     variables on your local machine or from those added through the
@@ -869,6 +1062,10 @@ def BigQuerySource(
         The path to the credentials file to use to connect.
     engine_args
         Additional arguments to use when constructing the SQLAlchemy engine.
+    temp_project
+        The BigQuery project to use for temporary tables.
+    temp_dataset
+        The BigQuery dataset to use for temporary tables.
 
     Returns
     -------
@@ -897,6 +1094,8 @@ def BigQuerySource(
     credentials_base64: Optional[str] = None,
     credentials_path: Optional[str] = None,
     engine_args: Optional[Dict[str, Any]] = None,
+    temp_project: Optional[str] = None,
+    temp_dataset: Optional[str] = None,
 ) -> BaseSQLSourceProtocol:
     """Create a BigQuery data source. SQL-based data sources
     created without arguments assume a configuration in your
@@ -914,6 +1113,8 @@ def BigQuerySource(
         location=location,
         credentials_base64=credentials_base64,
         credentials_path=credentials_path,
+        temp_project=temp_project,
+        temp_dataset=temp_dataset,
         engine_args=engine_args,
     )
 
@@ -1236,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
@@ -1256,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.
 
@@ -1274,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"),
+    ... )
     """
     ...
 
@@ -1286,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
@@ -1304,6 +1521,8 @@ def DatabricksSource(
         db=db,
         port=port,
         name=name,
+        client_id=client_id,
+        client_secret=client_secret,
         engine_args=engine_args,
     )
 
@@ -1747,6 +1966,8 @@ __all__ = (
     "DynamoDBSource",
     "FinalizedChalkQuery",
     "IncrementalSettings",
+    "MSSQLSource",
+    "MSSQLSourceImpl",
     "MySQLSource",
     "PostgreSQLSource",
     "RedshiftSource",

chalk/sql/_internal/integrations/athena.py

@@ -24,6 +24,7 @@ from chalk.sql.finalized_query import FinalizedChalkQuery
 from chalk.utils.df_utils import pa_array_to_pl_series
 from chalk.utils.log_with_context import get_logger
 from chalk.utils.missing_dependency import missing_dependency_exception
+from chalk.utils.pl_helpers import str_json_decode_compat
 from chalk.utils.threading import DEFAULT_IO_EXECUTOR, MultiSemaphore
 from chalk.utils.tracing import safe_incr, safe_set_gauge, safe_trace
 
@@ -485,7 +486,11 @@ class AthenaSourceImpl(BaseSQLSource):
                 if pa.types.is_list(expected_type) or pa.types.is_large_list(expected_type):
                     if pa.types.is_string(actual_type) or pa.types.is_large_string(actual_type):
                         series = pa_array_to_pl_series(tbl[col_name])
-                        column = series.str.json_extract(feature.converter.polars_dtype).to_arrow().cast(expected_type)
+                        column = (
+                            str_json_decode_compat(series, feature.converter.polars_dtype)
+                            .to_arrow()
+                            .cast(expected_type)
+                        )
                 if actual_type != expected_type:
                     column = column.cast(options=pc.CastOptions(target_type=expected_type, allow_time_truncate=True))
                 if isinstance(column, pa.ChunkedArray):

chalk/sql/_internal/integrations/bigquery.py

@@ -225,6 +225,8 @@ _BQ_DATASET_NAME = "BQ_DATASET"
 _BQ_PROJECT_NAME = "BQ_PROJECT"
 _BQ_CREDENTIALS_BASE64_NAME = "BQ_CREDENTIALS_BASE64"
 _BQ_CREDENTIALS_PATH_NAME = "BQ_CREDENTIALS_PATH"
+_BQ_TEMP_PROJECT_NAME = "BQ_TEMP_PROJECT"
+_BQ_TEMP_DATASET_NAME = "BQ_TEMP_DATASET"
 
 
 class BigQuerySourceImpl(BaseSQLSource):
@@ -239,6 +241,8 @@ class BigQuerySourceImpl(BaseSQLSource):
         location: Optional[str] = None,
         credentials_base64: Optional[str] = None,
         credentials_path: Optional[str] = None,
+        temp_project: Optional[str] = None,
+        temp_dataset: Optional[str] = None,
         engine_args: Optional[Dict[str, Any]] = None,
         integration_variable_override: Optional[Mapping[str, str]] = None,
     ):
@@ -267,6 +271,12 @@ class BigQuerySourceImpl(BaseSQLSource):
         self.credentials_path = credentials_path or load_integration_variable(
             integration_name=name, name=_BQ_CREDENTIALS_PATH_NAME, override=integration_variable_override
         )
+        self.temp_project = temp_project or load_integration_variable(
+            integration_name=name, name=_BQ_TEMP_PROJECT_NAME, override=integration_variable_override
+        )
+        self.temp_dataset = temp_dataset or load_integration_variable(
+            integration_name=name, name=_BQ_TEMP_DATASET_NAME, override=integration_variable_override
+        )
         BaseSQLSource.__init__(self, name=name, engine_args=engine_args, async_engine_args={})
 
     @functools.cached_property
@@ -397,6 +407,10 @@ class BigQuerySourceImpl(BaseSQLSource):
         except ModuleNotFoundError:
             raise missing_dependency_exception("chalkpy[bigquery]")
 
+        # Use temp_project/temp_dataset if specified, otherwise fall back to main project/dataset
+        temp_project = self.temp_project or self.project
+        temp_dataset = self.temp_dataset or self.dataset
+
         create_table_sql = create_temp_table.compile(dialect=self.get_sqlalchemy_dialect()).string
         create_table_sql = create_table_sql.replace("TEMPORARY", "", 1)
         chalk_logger.info(f"Creating temporary table {temp_table.name} in BigQuery {session_id=}: {create_table_sql}")
@@ -412,9 +426,10 @@ class BigQuerySourceImpl(BaseSQLSource):
             job_config=job_config,
         ).result()
         try:
+            temp_table_fqn = f"{temp_project}.{temp_dataset}.{temp_table.name}"
             connection.load_table_from_dataframe(
                 temp_value.to_pandas(),
-                f"{self.project}.{self.dataset}.{temp_table.name}",
+                temp_table_fqn,
                 job_config=google.cloud.bigquery.LoadJobConfig(connection_properties=connection_properties),
             ).result()
             yield
@@ -433,7 +448,10 @@ class BigQuerySourceImpl(BaseSQLSource):
     def _bigquery_output_table(self, client: google.cloud.bigquery.Client) -> Iterator[str]:
         destination_table_name = f"temp_output_{str(uuid4()).replace('-', '_')}"
 
-        destination = f"{client.project}.{self.dataset}.{destination_table_name}"
+        # Use temp_project/temp_dataset if specified, otherwise fall back to main project/dataset
+        temp_project = self.temp_project or self.project
+        temp_dataset = self.temp_dataset or self.dataset
+        destination = f"{temp_project}.{temp_dataset}.{destination_table_name}"
 
         try:
             yield destination
@@ -639,6 +657,8 @@ class BigQuerySourceImpl(BaseSQLSource):
                 create_integration_variable(_BQ_PROJECT_NAME, self.name, self.project),
                 create_integration_variable(_BQ_CREDENTIALS_BASE64_NAME, self.name, self.credentials_base64),
                 create_integration_variable(_BQ_CREDENTIALS_PATH_NAME, self.name, self.credentials_path),
+                create_integration_variable(_BQ_TEMP_PROJECT_NAME, self.name, self.temp_project),
+                create_integration_variable(_BQ_TEMP_DATASET_NAME, self.name, self.temp_dataset),
             ]
             if v is not None
         }

chalk/sql/_internal/integrations/databricks.py

@@ -17,6 +17,7 @@ from chalk.sql.finalized_query import FinalizedChalkQuery
 from chalk.utils.df_utils import pa_array_to_pl_series
 from chalk.utils.log_with_context import get_logger
 from chalk.utils.missing_dependency import missing_dependency_exception
+from chalk.utils.pl_helpers import str_json_decode_compat
 from chalk.utils.threading import DEFAULT_IO_EXECUTOR
 from chalk.utils.tracing import safe_incr, safe_trace
 
@@ -32,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):
@@ -45,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,
@@ -73,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:
@@ -96,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,
@@ -167,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")
 
@@ -226,7 +266,11 @@ class DatabricksSourceImpl(BaseSQLSource):
                 if pa.types.is_list(expected_type) or pa.types.is_large_list(expected_type):
                     if pa.types.is_string(actual_type) or pa.types.is_large_string(actual_type):
                         series = pa_array_to_pl_series(tbl[col_name])
-                        column = series.str.json_extract(feature.converter.polars_dtype).to_arrow().cast(expected_type)
+                        column = (
+                            str_json_decode_compat(series, feature.converter.polars_dtype)
+                            .to_arrow()
+                            .cast(expected_type)
+                        )
 
                 # Cast to expected type if needed
                 if actual_type != expected_type:
@@ -268,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)
@@ -348,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
         }