Support overriding ROWCOUNT for SQL Server (#19267)

* Support overriding ROWCOUNT for SQL Server

Add new "ROWCOUNT Override" connection property for `:sqlserver`, which will provide a DB-level mechanism to override the `ROWCOUNT` session level setting as needed for specific DBs

Change `max-results-bare-rows` from a hardcoded constant to a setting definition instead, which permits a DB level override, and move the former constant default to a new def instead (`default-max-results-bare-rows`)

For `:sqlserver`, set the DB-level setting override (if the connection property is set), via the `driver/normalize-db-details` impl

Add test to confirm the original scenario from #9940 works using this new override (set to `0`)

Move common computation function of overall row limit to the `metabase.query-processor.middleware.limit` namespace, and invoke it from execute now, called `determine-query-max-rows`

Add new clause to the `determine-query-max-rows` function that preferentially takes the value from `row-limit-override` (if defined)

bin/build-drivers/src/build_drivers/lint_manifest_file.clj

@@ -28,7 +28,7 @@
 (s/def ::parent (s/or :single-parent string? :multiple-parent (s/coll-of string?)))
 
 (s/def ::required boolean?)
-(s/def ::placeholder string?)
+(s/def ::placeholder any?)
 (s/def ::type #(contains? property-types %))
 (s/def ::visible-if (s/map-of keyword? any?))
 

modules/drivers/sqlserver/resources/metabase-plugin.yaml

@@ -23,6 +23,10 @@ driver:
     - password
     - cloud-ip-address-info
     - ssl
+    - name: rowcount-override
+      display-name: ROWCOUNT Override
+      placeholder: 0
+      required: false
     - ssh-tunnel
     - advanced-options-start
     - merge:

modules/drivers/sqlserver/src/metabase/driver/sqlserver.clj

@@ -447,6 +447,47 @@
   (let [parent-method (get-method sql.qp/preprocess :sql)]
     (fix-order-bys (parent-method driver inner-query))))
 
+;; In order to support certain native queries that might return results at the end, we have to use only prepared
+;; statements (see #9940)
+(defmethod sql-jdbc.execute/statement-supported? :sqlserver [_]
+  false)
+
+;; SQL server only supports setting holdability at the connection level, not the statement level, as per
+;; https://docs.microsoft.com/en-us/sql/connect/jdbc/using-holdability?view=sql-server-ver15
+;; and
+;; https://github.com/microsoft/mssql-jdbc/blob/v9.2.1/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerConnection.java#L5349-L5357
+;; an exception is thrown if they do not match, so it's safer to simply NOT try to override it at the statement level,
+;; because it's not supported anyway
+;; this impl is otherwise the same as the default
+(defmethod sql-jdbc.execute/prepared-statement :sqlserver
+  [driver ^Connection conn ^String sql params]
+  (let [stmt (.prepareStatement conn sql
+                                ResultSet/TYPE_FORWARD_ONLY
+                                ResultSet/CONCUR_READ_ONLY)]
+    (try
+      (.setFetchDirection stmt ResultSet/FETCH_FORWARD)
+      (sql-jdbc.execute/set-parameters! driver stmt params)
+      stmt
+      (catch Throwable e
+        (.close stmt)
+        (throw e)))))
+
+;; similar rationale to prepared-statement above
+(defmethod sql-jdbc.execute/statement :sqlserver
+  [_ ^Connection conn]
+  (let [stmt (.createStatement conn
+               ResultSet/TYPE_FORWARD_ONLY
+               ResultSet/CONCUR_READ_ONLY)]
+    (try
+      (try
+        (.setFetchDirection stmt ResultSet/FETCH_FORWARD)
+        (catch Throwable e
+          (log/debug e (trs "Error setting statement fetch direction to FETCH_FORWARD"))))
+      stmt
+      (catch Throwable e
+        (.close stmt)
+        (throw e)))))
+
 (defmethod unprepare/unprepare-value [:sqlserver LocalDate]
   [_ ^LocalDate t]
   ;; datefromparts(year, month, day)
@@ -506,41 +547,11 @@
   [driver bool]
   (sql/->prepared-substitution driver (if bool 1 0)))
 
-;; SQL server only supports setting holdability at the connection level, not the statement level, as per
-;; https://docs.microsoft.com/en-us/sql/connect/jdbc/using-holdability?view=sql-server-ver15
-;; and
-;; https://github.com/microsoft/mssql-jdbc/blob/v9.2.1/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerConnection.java#L5349-L5357
-;; an exception is thrown if they do not match, so it's safer to simply NOT try to override it at the statement level,
-;; because it's not supported anyway
-(defmethod sql-jdbc.execute/prepared-statement :sqlserver
-  [driver ^Connection conn ^String sql params]
-  (let [stmt (.prepareStatement conn
-                                sql
-                                ResultSet/TYPE_FORWARD_ONLY
-                                ResultSet/CONCUR_READ_ONLY)]
-    (try
-      (try
-        (.setFetchDirection stmt ResultSet/FETCH_FORWARD)
-        (catch Throwable e
-          (log/debug e (trs "Error setting prepared statement fetch direction to FETCH_FORWARD"))))
-      (sql-jdbc.execute/set-parameters! driver stmt params)
-      stmt
-      (catch Throwable e
-        (.close stmt)
-        (throw e)))))
-
-;; similar rationale to prepared-statement above
-(defmethod sql-jdbc.execute/statement :sqlserver
-  [_ ^Connection conn]
-  (let [stmt (.createStatement conn
-                               ResultSet/TYPE_FORWARD_ONLY
-                               ResultSet/CONCUR_READ_ONLY)]
-    (try
-      (try
-        (.setFetchDirection stmt ResultSet/FETCH_FORWARD)
-        (catch Throwable e
-          (log/debug e (trs "Error setting statement fetch direction to FETCH_FORWARD"))))
-      stmt
-      (catch Throwable e
-        (.close stmt)
-        (throw e)))))
+(defmethod driver/normalize-db-details :sqlserver
+  [_ database]
+  (if-let [rowcount-override (-> database :details :rowcount-override)]
+    ;; if the user has set the rowcount-override connection property, it ends up in the details map, but it actually
+    ;; needs to be moved over to the settings map (which is where DB local settings go, as per #19399)
+    (-> (update database :details #(dissoc % :rowcount-override))
+        (update :settings #(assoc % :max-results-bare-rows rowcount-override)))
+    database))

modules/drivers/sqlserver/test/metabase/driver/sqlserver_test.clj

@@ -12,8 +12,10 @@
             [metabase.driver.sql.query-processor :as sql.qp]
             [metabase.driver.sql.util.unprepare :as unprepare]
             [metabase.driver.sqlserver :as sqlserver]
+            [metabase.models :refer [Database]]
             [metabase.query-processor :as qp]
             [metabase.query-processor.interface :as qp.i]
+            [metabase.query-processor.middleware.constraints :as constraints]
             [metabase.query-processor.timezone :as qp.timezone]
             [metabase.test :as mt]))
 
@@ -296,3 +298,45 @@
                                (mt/run-mbql-query checkins
                                  {:aggregation [[:count]]
                                   :breakout    [[:field $date {:temporal-unit unit}]]}))))))))))))
+
+(deftest max-results-bare-rows-test
+  (mt/test-driver :sqlserver
+    (testing "Should support overriding the ROWCOUNT for a specific SQL Server DB (#9940)"
+      (mt/with-temp Database [db {:name    "SQL Server with ROWCOUNT override"
+                                  :engine  "sqlserver"
+                                  :details (-> (:details (mt/db))
+                                               ;; SQL server considers a ROWCOUNT of 0 to be unconstrained
+                                               ;; we are putting this in the details map, since that's where connection
+                                               ;; properties go in a client save operation, but it will be MOVED to the
+                                               ;; settings map instead (which is where DB-local settings go), via the
+                                               ;; driver/normalize-db-details implementation for :sqlserver
+                                               (assoc :rowcount-override 0))}]
+        (mt/with-db db
+          (is (= 3000 (-> {:query (str "DECLARE @DATA AS TABLE(\n"
+                                       "    IDX INT IDENTITY(1,1),\n"
+                                       "    V INT\n"
+                                       ")\n"
+                                       "DECLARE @STEP INT \n"
+                                       "SET @STEP = 1\n"
+                                       "WHILE @STEP <=3000\n"
+                                       "BEGIN\n"
+                                       "    INSERT INTO @DATA(V)\n"
+                                       "    SELECT 1\n"
+                                       "    SET @STEP = @STEP + 1\n"
+                                       "END \n"
+                                       "\n"
+                                       "DECLARE @TEMP AS TABLE(\n"
+                                       "    IDX INT IDENTITY(1,1),\n"
+                                       "    V INT\n"
+                                       ")\n"
+                                       "INSERT INTO @TEMP(V)\n"
+                                       "SELECT V FROM @DATA\n"
+                                       "\n"
+                                       "SELECT COUNT(1) FROM @TEMP\n")}
+                         mt/native-query
+                         ;; add default query constraints to ensure the default limit of 2000 is overridden by the
+                         ;; `:rowcount-override` connection property we defined in the details above
+                         (assoc :constraints constraints/default-query-constraints)
+                         qp/process-query
+                         mt/rows
+                         ffirst))))))))

src/metabase/driver/sql_jdbc/execute.clj

@@ -13,11 +13,10 @@
             [metabase.driver.sql-jdbc.execute.diagnostic :as sql-jdbc.execute.diagnostic]
             [metabase.driver.sql-jdbc.execute.old-impl :as execute.old]
             [metabase.driver.sql-jdbc.sync.interface :as sql-jdbc.sync]
-            [metabase.mbql.util :as mbql.u]
             [metabase.models.setting :refer [defsetting]]
             [metabase.query-processor.context :as context]
             [metabase.query-processor.error-type :as qp.error-type]
-            [metabase.query-processor.interface :as qp.i]
+            [metabase.query-processor.middleware.limit :as limit]
             [metabase.query-processor.reducible :as qp.reducible]
             [metabase.query-processor.store :as qp.store]
             [metabase.query-processor.timezone :as qp.timezone]
@@ -491,8 +490,7 @@
    {:pre [(string? sql) (seq sql)]}
    (let [remark   (qputil/query->remark driver outer-query)
          sql      (str "-- " remark "\n" sql)
-         max-rows (or (mbql.u/query->max-rows-limit outer-query)
-                      qp.i/absolute-max-results)]
+         max-rows (limit/determine-query-max-rows outer-query)]
      (execute-reducible-query driver sql params max-rows context respond)))
 
   ([driver sql params max-rows context respond]

src/metabase/driver/util.clj

@@ -194,8 +194,8 @@
       :description (trs "Comma separated names of {0} that <strong>should NOT</strong> appear in Metabase" (str/lower-case disp-name))
       :visible-if  {(keyword type-prop-nm) "exclusion"}
       :helper-text (trs "You can use patterns like <strong>auth*</strong> to match multiple {0}" (str/lower-case disp-name))
-      :required true}
-      ]))
+      :required true}]))
+
 
 (defn find-schema-filters-prop
   "Finds the first property of type `:schema-filters` for the given `driver` connection properties. Returns `nil`

src/metabase/query_processor/middleware/constraints.clj

 (ns metabase.query-processor.middleware.constraints
   "Middleware that adds default constraints to limit the maximum number of rows returned to queries that specify the
-  `:add-default-userland-constraints?` `:middleware` option.")
+  `:add-default-userland-constraints?` `:middleware` option."
+  (:require [metabase.models.setting :as setting]
+            [metabase.util.i18n :refer [deferred-tru]]))
 
-(def ^:private max-results-bare-rows
-  "Maximum number of rows to return specifically on :rows type queries via the API."
-  2000)
+(def ^:private ^:const default-max-results-bare-rows 2000)
+
+;; NOTE: this was changed from a hardcoded var with value of 2000 (now moved to `default-max-results-bare-rows`)
+;; to a setting in 0.43 the setting, which allows for DB local value, can still be nil, so any places below that used
+;; to reference the former constant value have to expect it could return nil instead
+(setting/defsetting max-results-bare-rows
+  (deferred-tru "Maximum number of rows to return specifically on :rows type queries via the API.")
+  :visibility     :authenticated
+  :type           :integer
+  :database-local :allowed)
 
 (def ^:private max-results
   "General maximum number of rows to return from an API query."
@@ -13,7 +22,7 @@
 (def default-query-constraints
   "Default map of constraints that we apply on dataset queries executed by the api."
   {:max-results           max-results
-   :max-results-bare-rows max-results-bare-rows})
+   :max-results-bare-rows (or (max-results-bare-rows) default-max-results-bare-rows)})
 
 (defn- ensure-valid-constraints
   "`:max-results-bare-rows` must be less than or equal to `:max-results`, so if someone sets `:max-results` but not

src/metabase/query_processor/middleware/limit.clj

 (ns metabase.query-processor.middleware.limit
   "Middleware that handles limiting the maximum number of rows returned by a query."
   (:require [metabase.mbql.util :as mbql.u]
+            [metabase.models.setting :as setting]
             [metabase.query-processor.interface :as i]
             [metabase.query-processor.util :as qputil]))
 
@@ -27,13 +28,25 @@
            (ensure-reduced result')
            result'))))))
 
+(defn determine-query-max-rows
+  "Given a `query`, return the max rows that should be returned.  This is the first non-nil value from (in decreasing
+  priority order):
+
+  1. the value of the `metabase.query-processor.middleware.constraints/max-results-bare-rows` setting, which allows
+     for database-local override
+  2. the output of `metabase.mbql.util/query->max-rows-limit` when called on the given query
+  3. `metabase.query-processor.interface/absolute-max-results` (a constant, non-nil backstop value)"
+  [query]
+  (or (setting/get-value-of-type :integer :max-results-bare-rows)
+      (mbql.u/query->max-rows-limit query)
+      i/absolute-max-results))
+
 (defn limit
   "Add an implicit `limit` clause to MBQL queries without any aggregations, and limit the maximum number of rows that
   can be returned in post-processing."
   [qp]
   (fn [query rff context]
-    (let [max-rows (or (mbql.u/query->max-rows-limit query)
-                       i/absolute-max-results)]
+    (let [max-rows (determine-query-max-rows query)]
       (qp
        (add-limit max-rows query)
        (fn [metadata]

test/metabase/query_processor/middleware/constraints_test.clj

@@ -15,7 +15,7 @@
   (testing "if it is *truthy* add the constraints"
     (is (= {:middleware  {:add-default-userland-constraints? true},
             :constraints {:max-results           @#'constraints/max-results
-                          :max-results-bare-rows @#'constraints/max-results-bare-rows}}
+                          :max-results-bare-rows @#'constraints/default-max-results-bare-rows}}
            (add-default-userland-constraints
             {:middleware {:add-default-userland-constraints? true}})))))