mage-ai 0.9.32__py3-none-any.whl → 0.9.34__py3-none-any.whl

mage_ai/data_preparation/models/block/__init__.py

@@ -551,6 +551,7 @@ class Block(DataIntegrationMixin):
         from mage_ai.data_preparation.models.block.sql.utils.shared import (
             extract_create_statement_table_name,
             extract_insert_statement_table_names,
+            extract_update_statement_table_names,
         )
 
         if not self.content:
@@ -561,6 +562,9 @@ class Block(DataIntegrationMixin):
             return table_name
 
         matches = extract_insert_statement_table_names(self.content)
+        if len(matches) == 0:
+            matches = extract_update_statement_table_names(self.content)
+
         if len(matches) == 0:
             return None
 
@@ -1036,7 +1040,7 @@ class Block(DataIntegrationMixin):
                         variable_mapping,
                         execution_partition=execution_partition,
                         override_outputs=True,
-                        spark=(global_vars or dict()).get('spark'),
+                        spark=self.__get_spark_session_from_global_vars(global_vars=global_vars),
                         dynamic_block_uuid=dynamic_block_uuid,
                     )
                 except ValueError as e:
@@ -1612,6 +1616,34 @@ class Block(DataIntegrationMixin):
             variable_uuid=variable_uuid,
         )
 
+    def get_raw_outputs(
+        self,
+        block_uuid: str,
+        execution_partition: str = None,
+        from_notebook: bool = False,
+        global_vars: Dict = None,
+    ) -> List[Any]:
+        all_variables = self.get_variables_by_block(
+            block_uuid=block_uuid,
+            partition=execution_partition,
+        )
+
+        outputs = []
+
+        for variable_uuid in all_variables:
+            variable = self.pipeline.get_block_variable(
+                block_uuid,
+                variable_uuid,
+                from_notebook=from_notebook,
+                global_vars=global_vars,
+                partition=execution_partition,
+                raise_exception=True,
+                spark=self.__get_spark_session_from_global_vars(global_vars),
+            )
+            outputs.append(variable)
+
+        return outputs
+
     def get_outputs(
         self,
         execution_partition: str = None,
@@ -2237,21 +2269,12 @@ df = get_variable('{self.pipeline.uuid}', '{block_uuid}', 'df')
         if not test_functions or len(test_functions) == 0:
             return
 
-        variable_manager = self.pipeline.variable_manager
-        outputs = [
-            variable_manager.get_variable(
-                self.pipeline.uuid,
-                self.uuid,
-                variable,
-                partition=execution_partition,
-                spark=(global_vars or dict()).get('spark'),
-            )
-            for variable in self.output_variables(
-                execution_partition=execution_partition,
-                from_notebook=from_notebook,
-                global_vars=global_vars,
-            )
-        ]
+        outputs = self.get_raw_outputs(
+            dynamic_block_uuid or self.uuid,
+            execution_partition=execution_partition,
+            from_notebook=from_notebook,
+            global_vars=global_vars,
+        )
 
         if logger and 'logger' not in global_vars:
             global_vars['logger'] = logger
@@ -2468,6 +2491,23 @@ df = get_variable('{self.pipeline.uuid}', '{block_uuid}', 'df')
         self.spark_init = True
         return self.spark
 
+    def __get_spark_session_from_global_vars(self, global_vars: Dict = None):
+        if global_vars is None:
+            global_vars = dict()
+        spark = global_vars.get('spark')
+        if self.pipeline and self.pipeline.spark_config:
+            spark_config = self.pipeline.spark_config
+        else:
+            repo_config = RepoConfig(repo_path=self.repo_path)
+            spark_config = repo_config.spark_config
+        if not spark_config:
+            return spark
+        spark_config = SparkConfig.load(config=spark_config)
+        if spark_config.use_custom_session:
+            return global_vars.get('context', dict()).get(
+                spark_config.custom_session_var_name, spark)
+        return spark
+
     def __store_variables_prepare(
         self,
         variable_mapping: Dict,

mage_ai/data_preparation/models/block/data_integration/mixins.py

@@ -53,6 +53,12 @@ class DataIntegrationMixin:
 
     @property
     def inputs_only_uuids(self) -> List[str]:
+        """
+        Get a list of UUIDs of input streams that are marked as 'input_only'.
+
+        Returns:
+            List[str]: List of UUIDs.
+        """
         arr = []
 
         for stream_id, settings in self.data_integration_inputs.items():
@@ -63,6 +69,12 @@ class DataIntegrationMixin:
 
     @property
     def data_integration_inputs(self) -> Dict:
+        """
+        Get the data integration inputs configuration for this block.
+
+        Returns:
+            Dict: Data integration inputs configuration.
+        """
         mapping = {}
 
         if self.configuration_data_integration:
@@ -77,6 +89,12 @@ class DataIntegrationMixin:
 
     @property
     def uuids_for_inputs(self) -> List[str]:
+        """
+        Get a list of UUIDs associated with data integration inputs.
+
+        Returns:
+            List[str]: List of UUIDs.
+        """
         arr = []
 
         for stream_id, settings in self.data_integration_inputs.items():
@@ -87,12 +105,27 @@ class DataIntegrationMixin:
 
     @property
     def upstream_block_uuids_for_inputs(self) -> List[str]:
+        """
+        Get a list of upstream block UUIDs associated with data integration inputs.
+
+        Returns:
+            List[str]: List of upstream block UUIDs.
+        """
         if self.configuration_data_integration:
             inputs_combined = self.uuids_for_inputs + self.inputs_only_uuids
 
             return [up_uuid for up_uuid in self.upstream_block_uuids if up_uuid in inputs_combined]
 
     def get_block_data_integration_settings_directory_path(self, block_uuid: str = None) -> str:
+        """
+        Get the directory path for storing data integration settings associated with a block.
+
+        Args:
+            block_uuid (str, optional): UUID of the block. Defaults to None.
+
+        Returns:
+            str: Directory path for data integration settings.
+        """
         if not self.pipeline:
             return
 
@@ -104,12 +137,27 @@ class DataIntegrationMixin:
         )
 
     def get_catalog_file_path(self, block_uuid: str = None) -> str:
+        """
+        Get the file path for the data integration catalog file associated with a block.
+
+        Args:
+            block_uuid (str, optional): UUID of the block. Defaults to None.
+
+        Returns:
+            str: File path for the catalog file.
+        """
         return os.path.join(
             self.get_block_data_integration_settings_directory_path(block_uuid),
             BLOCK_CATALOG_FILENAME,
         )
 
     def get_catalog_from_file(self) -> Dict:
+        """
+        Read and return the data integration catalog from a file.
+
+        Returns:
+            Dict: Data integration catalog.
+        """
         catalog_full_path = self.get_catalog_file_path()
 
         if os.path.exists(catalog_full_path):
@@ -122,6 +170,12 @@ class DataIntegrationMixin:
                         return
 
     async def get_catalog_from_file_async(self) -> Dict:
+        """
+        Asynchronously read and return the data integration catalog from a file.
+
+        Returns:
+            Dict: Data integration catalog.
+        """
         catalog_full_path = self.get_catalog_file_path()
 
         if os.path.exists(catalog_full_path):
@@ -133,7 +187,13 @@ class DataIntegrationMixin:
                     except json.decoder.JSONDecodeError:
                         return
 
-    def update_catalog_file(self, catalog: Dict = None) -> Dict:
+    def update_catalog_file(self, catalog: Dict = None) -> None:
+        """
+        Update the data integration catalog file with the provided catalog.
+
+        Args:
+            catalog (Dict, optional): Data integration catalog to be saved. Defaults to None.
+        """
         catalog_full_path = self.get_catalog_file_path()
 
         os.makedirs(os.path.dirname(catalog_full_path), exist_ok=True)
@@ -143,6 +203,13 @@ class DataIntegrationMixin:
                 f.write(json.dumps(catalog))
 
     def is_data_integration(self) -> bool:
+        """
+        Check if the block is a data integration block.
+        If the data_integration_in_batch_pipeline feature is not enabled, return False.
+
+        Returns:
+            bool: True if it's a data integration block, False otherwise.
+        """
         if not self.pipeline or not \
                 Project(self.pipeline.repo_config).is_feature_enabled(
                     FeatureUUID.DATA_INTEGRATION_IN_BATCH_PIPELINE,
@@ -185,6 +252,30 @@ class DataIntegrationMixin:
         partition: str = None,
         **kwargs,
     ) -> Dict:
+        """
+        Retrieve data integration settings for the current block.
+
+        Args:
+            data_integration_uuid_only (bool, optional): If True, retrieve only the data integration
+                UUID. Defaults to False.
+            dynamic_block_index (Union[int, None], optional): The index of the dynamic block, if
+                applicable. Defaults to None.
+            dynamic_upstream_block_uuids (Union[List[str], None], optional): List of upstream block
+                UUIDs, if applicable. Defaults to None.
+            from_notebook (bool, optional): Whether the request is made from a notebook context.
+                Defaults to False.
+            global_vars (Dict, optional): Global variables to be used in template rendering.
+                Defaults to None.
+            input_vars (List[Any], optional): Input variables for the block. Defaults to None.
+            partition (str, optional): Partition identifier, if applicable.
+                Defaults to None.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            Dict: A dictionary containing the data integration settings, including catalog, config,
+            data_integration_uuid, selected_streams, and the appropriate key
+            (source or destination).
+        """
         if not self.is_data_integration():
             return
 

mage_ai/data_preparation/models/block/data_integration/utils.py

@@ -13,6 +13,7 @@ from mage_ai.data_integrations.logger.utils import (
     print_log_from_line,
     print_logs_from_output,
 )
+from mage_ai.data_integrations.utils.config import get_batch_fetch_limit
 from mage_ai.data_integrations.utils.parsers import parse_logs_and_json
 from mage_ai.data_preparation.models.block.data_integration.constants import (
     EXECUTION_PARTITION_FROM_NOTEBOOK,
@@ -383,21 +384,20 @@ def execute_data_integration(
         # Handle incremental sync
         state_file_path = None
         if index is not None:
+            batch_fetch_limit = get_batch_fetch_limit(config)
             state_file_path = get_state_file_path(block, data_integration_uuid, stream)
-
             stream_catalogs = get_streams_from_catalog(catalog, [stream]) or []
+
             if len(stream_catalogs) == 1 and \
                     REPLICATION_METHOD_INCREMENTAL == stream_catalogs[0].get('replication_method'):
                 # Use the state to adjust the query
                 # How do we write to the state when the source syncs can run in parallel?
                 pass
             else:
-                from mage_integrations.sources.constants import BATCH_FETCH_LIMIT
-                query_data['_offset'] = BATCH_FETCH_LIMIT * index
+                query_data['_offset'] = batch_fetch_limit * index
 
             if not is_last_block_run:
-                from mage_integrations.sources.constants import BATCH_FETCH_LIMIT
-                query_data['_limit'] = BATCH_FETCH_LIMIT
+                query_data['_limit'] = batch_fetch_limit
 
         tags = dict(block_tags=dict(
             index=index,

mage_ai/data_preparation/models/block/integration/__init__.py

@@ -10,7 +10,11 @@ import pandas as pd
 
 from mage_ai.data_cleaner.transformer_actions.utils import clean_column_name
 from mage_ai.data_integrations.logger.utils import print_log_from_line
-from mage_ai.data_integrations.utils.config import build_config, get_catalog_by_stream
+from mage_ai.data_integrations.utils.config import (
+    build_config,
+    get_batch_fetch_limit,
+    get_catalog_by_stream,
+)
 from mage_ai.data_preparation.models.block import Block
 from mage_ai.data_preparation.models.constants import PYTHON_COMMAND, BlockType
 from mage_ai.data_preparation.shared.stream import StreamToLogger
@@ -50,8 +54,6 @@ class IntegrationBlock(Block):
         runtime_arguments: Dict = None,
         **kwargs,
     ) -> List:
-        from mage_integrations.sources.constants import BATCH_FETCH_LIMIT
-
         if logging_tags is None:
             logging_tags = dict()
 
@@ -106,10 +108,6 @@ class IntegrationBlock(Block):
                     source_state_file_path,
                     destination_state_file_path,
                 )
-            else:
-                query_data['_offset'] = BATCH_FETCH_LIMIT * index
-            if not is_last_block_run:
-                query_data['_limit'] = BATCH_FETCH_LIMIT
 
         outputs = []
         if BlockType.DATA_LOADER == self.type:
@@ -120,6 +118,13 @@ class IntegrationBlock(Block):
                     self.pipeline.data_loader.file_path,
                     variables_dictionary_for_config,
                 )
+                batch_fetch_limit = get_batch_fetch_limit(config)
+
+                if stream_catalog.get('replication_method') != 'INCREMENTAL':
+                    query_data['_offset'] = batch_fetch_limit * index
+                if not is_last_block_run:
+                    query_data['_limit'] = batch_fetch_limit
+
                 args = [
                     PYTHON_COMMAND,
                     self.pipeline.source_file_path,

mage_ai/data_preparation/models/block/sql/__init__.py

@@ -19,6 +19,7 @@ from mage_ai.data_preparation.models.block.sql import (
 from mage_ai.data_preparation.models.block.sql.utils.shared import (
     has_create_or_insert_statement,
     has_drop_statement,
+    has_update_statement,
     interpolate_vars,
     split_query_string,
     table_name_parts_from_query,
@@ -29,6 +30,7 @@ from mage_ai.io.config import ConfigFileLoader
 from mage_ai.settings.repo import get_repo_path
 
 PREVIEWABLE_BLOCK_TYPES = [
+    BlockType.CUSTOM,
     BlockType.DATA_EXPORTER,
     BlockType.DATA_LOADER,
     BlockType.DBT,
@@ -646,16 +648,17 @@ def execute_raw_sql(
 
     has_create_or_insert = has_create_or_insert_statement(query_string)
     has_drop = has_drop_statement(query_string)
+    has_update = has_update_statement(query_string)
 
     for query in split_query_string(query_string):
-        if has_create_or_insert or has_drop:
+        if has_create_or_insert or has_drop or has_update:
             queries.append(query)
             fetch_query_at_indexes.append(False)
         else:
             queries.append(query)
             fetch_query_at_indexes.append(True)
 
-    if should_query and has_create_or_insert:
+    if should_query and (has_create_or_insert or has_update) and block.full_table_name:
         queries.append(f'SELECT * FROM {block.full_table_name} LIMIT 1000')
         fetch_query_at_indexes.append(block.full_table_name)
 

mage_ai/data_preparation/models/block/sql/utils/shared.py

@@ -494,7 +494,16 @@ def extract_insert_statement_table_names(text: str) -> List[str]:
 
 def extract_drop_statement_table_names(text: str) -> List[str]:
     matches = re.findall(
-        r'drop table(?: if exists)*',
+        r'\bdrop\s+table(?:\s+if\s+exists)?\s+([\w.]+)',
+        remove_comments(text),
+        re.IGNORECASE,
+    )
+    return matches
+
+
+def extract_update_statement_table_names(text: str) -> List[str]:
+    matches = re.findall(
+        r'\bupdate\b\s+([\w.]+)\s+set\s+[\s\S]*?\bwhere\b',
         remove_comments(text),
         re.IGNORECASE,
     )
@@ -515,6 +524,11 @@ def has_drop_statement(text: str) -> bool:
     return len(matches) >= 1
 
 
+def has_update_statement(text: str) -> bool:
+    matches = extract_update_statement_table_names(text)
+    return len(matches) >= 1
+
+
 def split_query_string(query_string: str) -> List[str]:
     text_parts = []
 

mage_ai/data_preparation/models/pipelines/interactions.py

@@ -326,6 +326,21 @@ class PipelineInteractions:
 
         await self.save()
 
+    async def variables(self) -> Dict:
+        variables = {}
+
+        interaction = await self.interaction()
+        if not interaction:
+            return variables
+
+        for block_interactions in (interaction.blocks or {}).values():
+            for block_interaction in block_interactions:
+                interaction_base = InteractionBase(block_interaction.uuid)
+                interaction_variables = await interaction_base.variables()
+                variables.update(interaction_variables)
+
+        return variables
+
     async def to_dict(self) -> Dict:
         interaction = await self.interaction()
 

mage_ai/data_preparation/shared/secrets.py

@@ -155,6 +155,7 @@ def get_secret_value(
     pipeline_uuid: str = None,
     project_uuid: str = None,
     repo_name: str = None,
+    **kwargs,
 ) -> Optional[str]:
     from mage_ai.orchestration.db.models.secrets import Secret
     key, key_uuid = _get_encryption_key(
@@ -189,29 +190,17 @@ def get_secret_value(
                     return fernet.decrypt(secret_legacy.value.encode('utf-8')).decode('utf-8')
                 except InvalidToken:
                     pass
+    if not kwargs.get('suppress_warning', False):
+        print(f'WARNING: Could not find secret value for secret {name}.')
 
-    print(f'WARNING: Could not find secret value for secret {name}.')
 
-
-def get_secret_value_db_safe(
-    name: str,
-    entity: Entity = Entity.GLOBAL,
-    pipeline_uuid: str = None,
-    project_uuid: str = None,
-    repo_name: str = None,
-) -> Optional[str]:
+def get_secret_value_db_safe(name: str, **kwargs) -> Optional[str]:
     """
     Calls get_secret_value only if the db has already been initialized.
     """
     from mage_ai.orchestration.db import db_connection
     if db_connection.session and db_connection.session.is_active:
-        return get_secret_value(
-            name,
-            entity=entity,
-            pipeline_uuid=pipeline_uuid,
-            project_uuid=project_uuid,
-            repo_name=repo_name,
-        )
+        return get_secret_value(name, **kwargs)
     else:
         return None
 
@@ -222,6 +211,7 @@ def delete_secret(
     entity: Entity = Entity.GLOBAL,
     pipeline_uuid: str = None,
     project_uuid: str = None,
+    **kwargs,
 ) -> None:
     from mage_ai.orchestration.db.models.secrets import Secret
     secret = None
@@ -242,7 +232,7 @@ def delete_secret(
 
     if secret:
         secret.delete()
-    else:
+    elif not kwargs.get('suppress_warning', False):
         print(f'WARNING: Could not find secret {name}')
 
 

mage_ai/data_preparation/templates/data_exporters/streaming/google_cloud_pubsub.yaml

@@ -0,0 +1,5 @@
+connector_type: google_cloud_pubsub
+project_id: my-project-id
+topic_id: my-topic-id
+pubsub_emulator_host: null  # "host.docker.internal:8085"
+path_to_credentials_json_file: "./google_credentials.json"

mage_ai/data_preparation/templates/repo/metadata.yaml

@@ -37,6 +37,11 @@ spark_config:
   # List of key-value pairs to be set in SparkConf
   # e.g., others: {'spark.executor.memory': '4g', 'spark.executor.cores': '2'}
   others: {}
+  # Wehther to create custom SparkSession via code and set in kwargs['context']
+  use_custom_session: false
+  # The variable name to set in kwargs['context'],
+  # e.g. kwargs['context']['spark'] = spark_session
+  custom_session_var_name: 'spark'
 
 notification_config:
   alert_on:

mage_ai/io/druid.py

@@ -1,9 +1,12 @@
-from mage_ai.io.base import ExportWritePolicy, QUERY_ROW_LIMIT
+from typing import IO, Any, Dict, List, Union
+
+from pandas import DataFrame, Series
+from pydruid.db.api import Connection
+from pydruid.db.api import Cursor as CursorParent
+
+from mage_ai.io.base import QUERY_ROW_LIMIT, ExportWritePolicy
 from mage_ai.io.config import BaseConfigLoader, ConfigKey
 from mage_ai.io.sql import BaseSQL
-from pandas import DataFrame, Series
-from pydruid.db.api import Connection, Cursor as CursorParent
-from typing import Dict, List, Union, Any, IO
 
 WRITE_NOT_SUPPORTED_EXCEPTION = Exception('write operations are not supported.')
 
@@ -138,7 +141,8 @@ class Druid(BaseSQL):
                          db_dtypes: List[str],
                          dtypes: List[str],
                          full_table_name: str,
-                         buffer: Union[IO, None] = None) -> None:
+                         buffer: Union[IO, None] = None,
+                         **kwargs) -> None:
         raise WRITE_NOT_SUPPORTED_EXCEPTION
 
     def export(self,

mage_ai/io/mssql.py

@@ -4,8 +4,9 @@ import numpy as np
 import pyodbc
 import simplejson
 from pandas import DataFrame, Series
+from sqlalchemy import create_engine
 
-from mage_ai.io.base import QUERY_ROW_LIMIT
+from mage_ai.io.base import QUERY_ROW_LIMIT, ExportWritePolicy
 from mage_ai.io.config import BaseConfigLoader, ConfigKey
 from mage_ai.io.export_utils import PandasTypes
 from mage_ai.io.sql import BaseSQL
@@ -21,6 +22,7 @@ class MSSQL(BaseSQL):
         user: str,
         schema: str = None,
         port: int = 1433,
+        verbose: bool = True,
         **kwargs,
     ) -> None:
         super().__init__(
@@ -30,24 +32,25 @@ class MSSQL(BaseSQL):
             password=password,
             schema=schema,
             port=port,
+            verbose=verbose,
             **kwargs
         )
 
-    def _enforce_limit(self, query: str, limit: int = QUERY_ROW_LIMIT) -> str:
-        # MSSQL doesn't support WITH statements in subqueries, so if the user uses a WITH
-        # statement in their query, it would break if we tried to enforce a limit.
-        return query
-
-    @classmethod
-    def with_config(cls, config: BaseConfigLoader) -> 'MSSQL':
-        return cls(
-            database=config[ConfigKey.MSSQL_DATABASE],
-            schema=config[ConfigKey.MSSQL_SCHEMA],
-            driver=config[ConfigKey.MSSQL_DRIVER],
-            host=config[ConfigKey.MSSQL_HOST],
-            password=config[ConfigKey.MSSQL_PASSWORD],
-            port=config[ConfigKey.MSSQL_PORT],
-            user=config[ConfigKey.MSSQL_USER],
+    @property
+    def connection_string(self):
+        driver = self.settings['driver']
+        server = self.settings['server']
+        database = self.settings['database']
+        username = self.settings['user']
+        password = self.settings['password']
+        return (
+            f'DRIVER={{{driver}}};'
+            f'SERVER={server};'
+            f'DATABASE={database};'
+            f'UID={username};'
+            f'PWD={password};'
+            'ENCRYPT=yes;'
+            'TrustServerCertificate=yes;'
         )
 
     def default_schema(self) -> str:
@@ -55,21 +58,9 @@ class MSSQL(BaseSQL):
 
     def open(self) -> None:
         with self.printer.print_msg('Opening connection to MSSQL database'):
-            driver = self.settings['driver']
-            server = self.settings['server']
-            database = self.settings['database']
-            username = self.settings['user']
-            password = self.settings['password']
-            connection_string = (
-                f'DRIVER={{{driver}}};'
-                f'SERVER={server};'
-                f'DATABASE={database};'
-                f'UID={username};'
-                f'PWD={password};'
-                'ENCRYPT=yes;'
-                'TrustServerCertificate=yes;'
+            self._ctx = pyodbc.connect(
+                self.connection_string,
             )
-            self._ctx = pyodbc.connect(connection_string)
 
     def build_create_schema_command(
         self,
@@ -149,6 +140,20 @@ class MSSQL(BaseSQL):
         sql = f'INSERT INTO {full_table_name} VALUES ({values_placeholder})'
         cursor.executemany(sql, values)
 
+    def upload_dataframe_fast(
+        self,
+        df: DataFrame,
+        schema_name: str,
+        table_name: str,
+        if_exists: ExportWritePolicy = ExportWritePolicy.REPLACE,
+
+    ):
+        engine = create_engine(
+            f'mssql+pyodbc://?odbc_connect={self.connection_string}',
+            fast_executemany=True,
+        )
+        df.to_sql(table_name, engine, schema=schema_name, if_exists=if_exists, index=False)
+
     def get_type(self, column: Series, dtype: str) -> str:
         if dtype in (
             PandasTypes.MIXED,
@@ -199,3 +204,20 @@ class MSSQL(BaseSQL):
             print(f'Invalid datatype provided: {dtype}')
 
         return 'char(255)'
+
+    @classmethod
+    def with_config(cls, config: BaseConfigLoader) -> 'MSSQL':
+        return cls(
+            database=config[ConfigKey.MSSQL_DATABASE],
+            schema=config[ConfigKey.MSSQL_SCHEMA],
+            driver=config[ConfigKey.MSSQL_DRIVER],
+            host=config[ConfigKey.MSSQL_HOST],
+            password=config[ConfigKey.MSSQL_PASSWORD],
+            port=config[ConfigKey.MSSQL_PORT],
+            user=config[ConfigKey.MSSQL_USER],
+        )
+
+    def _enforce_limit(self, query: str, limit: int = QUERY_ROW_LIMIT) -> str:
+        # MSSQL doesn't support WITH statements in subqueries, so if the user uses a WITH
+        # statement in their query, it would break if we tried to enforce a limit.
+        return query