(Fix #8267) Handle subscriptions in MSSQL when results exceed 2048 characters

PR-URL: https://github.com/hasura/graphql-engine-mono/pull/3959
GitOrigin-RevId: ea037c9dc4392d1b98ee086f6c87f79ce8ea5c8f

CHANGELOG.md

@@ -142,7 +142,8 @@ function:
 
 ### Bug fixes and improvements
 
-- server: add jsonb to string cast support - postgres
+- server: Fix regression in MSSQL subscriptions when results exceed 2048 characters (#8267)
+- server: add jsonb to string cast support - postgres 
 - server: improve performance of fetching postgres catalog metadata for tables and functions
 - server: Queries present in query collections, such as allow-list, and rest-endpoints are now validated (against the schema)
 - server: Redesigns internal implementation of webhook transforms.

server/src-lib/Database/MSSQL/Transaction.hs

@@ -11,6 +11,7 @@ module Database.MSSQL.Transaction
     singleRowQueryE,
     multiRowQuery,
     multiRowQueryE,
+    forJsonQueryE,
     buildGenericQueryTxE,
     withTxET,
   )
@@ -123,6 +124,24 @@ singleRowQueryE ef = rawQueryE ef singleRowResult
     singleRowResult (MSSQLResult [row]) = ODBC.fromRow row
     singleRowResult (MSSQLResult _) = Left "expecting single row"
 
+-- | MSSQL splits up results that have a @SELECT .. FOR JSON@ at the top-level
+-- into multiple rows with a single column, see
+-- https://docs.microsoft.com/en-us/sql/relational-databases/json/format-query-results-as-json-with-for-json-sql-server?view=sql-server-ver15#output-of-the-for-json-clause
+--
+-- This function simply concatenates each single-column row into one long 'Text' string.
+forJsonQueryE ::
+  forall m e.
+  MonadIO m =>
+  (MSSQLTxError -> e) ->
+  ODBC.Query ->
+  TxET e m Text
+forJsonQueryE ef = rawQueryE ef concatRowResult
+  where
+    concatRowResult :: MSSQLResult -> Either String Text
+    concatRowResult (MSSQLResult []) = pure mempty
+    concatRowResult (MSSQLResult rows@(r1 : _)) | length r1 == 1 = mconcat <$> mapM ODBC.fromRow rows
+    concatRowResult (MSSQLResult (r1 : _)) = Left $ "forJsonQueryE: Expected single-column results, but got " <> show (length r1) <> " columns"
+
 -- | Useful for building query transactions which return multiple rows.
 --
 -- @

server/src-lib/Hasura/Backends/MSSQL/Instances/Execute.hs

@@ -150,9 +150,7 @@ msDBSubscriptionExplain (SubscriptionQueryPlan plan sourceConfig variables _) =
   explainInfo <- liftEitherM $ runExceptT $ (mssqlRunReadOnly mssqlExecCtx) (runShowplan query)
   pure $ SubscriptionQueryPlanExplanation (T.toTxt query) explainInfo variables
 
---------------------------------------------------------------------------------
--- Producing the correct SQL-level list comprehension to multiplex a query
-
+-- | Producing the correct SQL-level list comprehension to multiplex a query
 -- Problem description:
 --
 -- Generate a query that repeats the same query N times but with
@@ -160,7 +158,9 @@ msDBSubscriptionExplain (SubscriptionQueryPlan plan sourceConfig variables _) =
 --
 -- [ Select x y | (x,y) <- [..] ]
 --
-
+-- Caution: Be aware that this query has a @FOR JSON@ clause at the top-level
+-- and hence its results may be split up across multiple rows. Use
+-- 'Database.MSSQL.Transaction.forJsonQueryE' to handle this.
 multiplexRootReselect ::
   [(CohortId, CohortVariables)] ->
   TSQL.Reselect ->

server/src-lib/Hasura/Backends/MSSQL/Instances/Transport.hs

@@ -12,7 +12,7 @@ import Data.ByteString qualified as B
 import Data.String (fromString)
 import Data.Text.Encoding (encodeUtf8)
 import Data.Text.Extended
-import Database.MSSQL.Transaction (singleRowQueryE)
+import Database.MSSQL.Transaction (forJsonQueryE)
 import Database.ODBC.SQLServer qualified as ODBC
 import Hasura.Backends.MSSQL.Connection
 import Hasura.Backends.MSSQL.Instances.Execute
@@ -121,7 +121,10 @@ executeMultiplexedQuery mssqlExecCtx query = do
   let parseResult r = J.eitherDecodeStrict (encodeUtf8 r) `onLeft` \s -> throw400 ParseFailed (fromString s)
       convertFromJSON :: [CohortResult] -> [(CohortId, B.ByteString)]
       convertFromJSON = map \(CohortResult (cid, cresult)) -> (cid, encodeUtf8 cresult)
-  textResult <- run $ mssqlRunReadOnly mssqlExecCtx $ singleRowQueryE defaultMSSQLTxErrorHandler query
+  -- Because the 'query' will have a @FOR JSON@ clause at the toplevel it will
+  -- be split across multiple rows, hence use of 'forJsonQueryE' which takes
+  -- care of concatenating the results.
+  textResult <- run $ mssqlRunReadOnly mssqlExecCtx $ forJsonQueryE defaultMSSQLTxErrorHandler query
   parsedResult <- parseResult textResult
   pure $ convertFromJSON parsedResult
 

server/tests-py/queries/subscriptions/basic/schema_teardown_mssql.yaml

@@ -5,3 +5,4 @@ args:
     source: mssql
     sql: |
       drop table hge_tests.test_t1
+      drop schema hge_tests

server/tests-py/queries/subscriptions/mssql/schema_setup_mssql.yaml

@@ -0,0 +1,61 @@
+type: bulk
+args:
+
+- type: mssql_run_sql
+  args:
+    source: mssql
+    sql: |
+      create schema hge_tests
+
+      create table hge_tests.test_subscriptions (
+        id int primary key,
+        field1 varchar(max)
+      )
+
+- type: mssql_run_sql
+  args:
+    source: mssql
+    sql: |
+      INSERT INTO
+        hge_tests.test_subscriptions(id, field1)
+        VALUES
+        (1, '1234567890'),
+        (2, '12345678901234567890'),
+        (3, '123456789012345678901234567890'),
+        (4, '1234567890123456789012345678901234567890'),
+        (5, '12345678901234567890123456789012345678901234567890'),
+        (6, '123456789012345678901234567890123456789012345678901234567890'),
+        (7, '1234567890123456789012345678901234567890123456789012345678901234567890'),
+        (8, '12345678901234567890123456789012345678901234567890123456789012345678901234567890'),
+        (9, '123456789012345678901234567890123456789012345678901234567890123456789012345678901234567890'),
+        (10, '1234567890123456789012345678901234567890123456789012345678901234567890123456789012345678901234567890'),
+        (11, '1234567890'),
+        (12, '12345678901234567890'),
+        (13, '123456789012345678901234567890'),
+        (14, '1234567890123456789012345678901234567890'),
+        (15, '12345678901234567890123456789012345678901234567890'),
+        (16, '123456789012345678901234567890123456789012345678901234567890'),
+        (17, '1234567890123456789012345678901234567890123456789012345678901234567890'),
+        (18, '12345678901234567890123456789012345678901234567890123456789012345678901234567890'),
+        (19, '123456789012345678901234567890123456789012345678901234567890123456789012345678901234567890'),
+        (20, '1234567890123456789012345678901234567890123456789012345678901234567890123456789012345678901234567890'),
+        (21, '1234567890'),
+        (22, '12345678901234567890'),
+        (23, '123456789012345678901234567890'),
+        (24, '1234567890123456789012345678901234567890'),
+        (25, '12345678901234567890123456789012345678901234567890'),
+        (26, '123456789012345678901234567890123456789012345678901234567890'),
+        (27, '1234567890123456789012345678901234567890123456789012345678901234567890'),
+        (28, '12345678901234567890123456789012345678901234567890123456789012345678901234567890'),
+        (29, '123456789012345678901234567890123456789012345678901234567890123456789012345678901234567890'),
+        (30, '1234567890123456789012345678901234567890123456789012345678901234567890123456789012345678901234567890'),
+        (31, '1234567890'),
+        (32, '12345678901234567890'),
+        (33, '123456789012345678901234567890'),
+        (34, '1234567890123456789012345678901234567890'),
+        (35, '12345678901234567890123456789012345678901234567890'),
+        (36, '123456789012345678901234567890123456789012345678901234567890'),
+        (37, '1234567890123456789012345678901234567890123456789012345678901234567890'),
+        (38, '12345678901234567890123456789012345678901234567890123456789012345678901234567890'),
+        (39, '123456789012345678901234567890123456789012345678901234567890123456789012345678901234567890'),
+        (40, '1234567890123456789012345678901234567890123456789012345678901234567890123456789012345678901234567890')

server/tests-py/queries/subscriptions/mssql/schema_teardown_mssql.yaml

@@ -0,0 +1,9 @@
+type: bulk
+args:
+- type: mssql_run_sql
+  args:
+    source: mssql
+    sql: |
+      drop table hge_tests.test_subscriptions
+
+      drop schema hge_tests

server/tests-py/queries/subscriptions/mssql/setup_mssql.yaml

@@ -0,0 +1,9 @@
+type: bulk
+args:
+
+- type: mssql_track_table
+  args:
+    source: mssql
+    table:
+      schema: hge_tests
+      name: test_subscriptions

server/tests-py/queries/subscriptions/mssql/teardown_mssql.yaml

@@ -0,0 +1,9 @@
+type: bulk
+args:
+
+- type: mssql_untrack_table
+  args:
+    source: mssql
+    table:
+      schema: hge_tests
+      name: test_subscriptions

server/tests-py/test_subscriptions.py

@@ -197,7 +197,7 @@ class TestSubscriptionBasic:
         ev = ws_client.get_ws_query_event('2',3)
         assert ev['type'] == 'complete' and ev['id'] == '2', ev
 
-## NOTE: The same tests as in TestSubcscriptionBasic but with
+## NOTE: The same tests as in TestSubscriptionBasic but with
 ##       the subscription transport being used is `graphql-ws`
 ## FIXME: There's an issue with the tests being parametrized with both
 ##        postgres and mssql data sources enabled(See issue #2084).
@@ -644,3 +644,31 @@ class TestSubscriptionCustomizedSourceCommon:
         ev = ws_client.get_ws_query_event('2',15)
         assert ev['type'] == 'error' and ev['id'] == '2', ev
         assert ev['payload']['errors'] == [OrderedDict([('extensions', OrderedDict([('path', '$'), ('code', 'validation-failed')])), ('message', 'subscriptions must select one top level field')])], ev
+
+@pytest.mark.parametrize("backend", ['mssql'])
+@usefixtures('per_class_tests_db_state', 'ws_conn_init')
+class TestSubscriptionMSSQLChunkedResults:
+    @classmethod
+    def dir(cls):
+        return 'queries/subscriptions/mssql'
+
+    query = """
+     subscription {
+       hge_tests_test_subscriptions {
+         field1
+       }
+     }
+    """
+
+    def test_chunked_results(self, ws_client):
+        obj = {
+            'id': '1',
+            'payload': {
+                'query': self.query
+            },
+            'type': 'start'
+        }
+        ws_client.send(obj)
+        ev = ws_client.get_ws_query_event('1',15)
+        assert ev['type'] == 'data' and ev['id'] == '1', ev
+        assert not "errors" in ev['payload'], ev