Tweak InputStream lifetimes

enterprise/backend/src/metabase_enterprise/transforms_python/execute.clj

@@ -11,12 +11,15 @@
    [metabase.driver :as driver]
    [metabase.driver.util :as driver.u]
    [metabase.util :as u]
+   [metabase.util.i18n :as i18n]
    [metabase.util.jvm :as u.jvm]
    [metabase.util.log :as log]
    [toucan2.core :as t2])
   (:import
-   (java.io Closeable File)
+   (java.io Closeable)
    (java.net SocketException)
+   (java.nio.file Files)
+   (java.nio.file.attribute FileAttribute)
    (java.time Duration)))
 
 (set! *warn-on-reflection* true)
@@ -243,13 +246,14 @@
   ;; TODO restructure things such that s3 can we swapped out for other transfer mechanisms
   (with-open [log-future-ref     (open-python-message-update-future! run-id message-log)
               shared-storage-ref (s3/open-shared-storage! (:source-tables source))]
-    (let [driver     (:engine db)
-          server-url (transforms-python.settings/python-runner-url)
-          _          (python-runner/copy-tables-to-s3! {:run-id         run-id
-                                                        :shared-storage @shared-storage-ref
-                                                        :table-name->id (:source-tables source)
-                                                        :cancel-chan    cancel-chan})
-          _          (start-cancellation-process! server-url run-id cancel-chan) ; inherits lifetime of cancel-chan
+    (let [driver          (:engine db)
+          server-url      (transforms-python.settings/python-runner-url)
+          _               (python-runner/copy-tables-to-s3! {:run-id         run-id
+                                                             :shared-storage @shared-storage-ref
+                                                             :table-name->id (:source-tables source)
+                                                             :cancel-chan    cancel-chan})
+          _               (start-cancellation-process! server-url run-id cancel-chan) ; inherits lifetime of cancel-chan
+
           {:keys [status body] :as response}
           (python-runner/execute-python-code-http-call!
            {:server-url     server-url
@@ -257,30 +261,28 @@
             :run-id         run-id
             :table-name->id (:source-tables source)
             :shared-storage @shared-storage-ref})
-          ;; TODO temporary to keep more code stable while refactoring
-          ;; no need to materialize these early (i.e output we can stream directly into a tmp file or db if small)
-          {:keys [output output-manifest events]} (python-runner/read-output-objects @shared-storage-ref)]
+
+          output-manifest (python-runner/read-output-manifest @shared-storage-ref)
+          events          (python-runner/read-events @shared-storage-ref)]
       (.close log-future-ref)                               ; early close to force any writes to flush
       (when (seq events)
         (replace-python-logs! message-log events))
       (if (not= 200 status)
         (throw (ex-info "Python runner call failed"
-                        {:status-code           400
-                         :api-status-code       status
-                         :body                  body
-                         :events                events
-                         :transform-run-message (message-log->transform-run-message message-log)}))
+                        {:transform-message (i18n/tru "Python execution failure (exit code {0})" (:exit_code body "?"))
+                         :status-code       400
+                         :api-status-code   status
+                         :body              body}))
         (try
-          (let [temp-file (File/createTempFile "transform-output-" ".jsonl")]
+          (let [temp-path (Files/createTempFile "transform-output-" ".jsonl" (u/varargs FileAttribute))
+                temp-file (.toFile temp-path)]
             (when-not (seq (:fields output-manifest))
               (throw (ex-info "No fields in metadata"
                               {:metadata               output-manifest
-                               :raw-body               body
-                               :events                 events
-                               :transform-run-message  (message-log->transform-run-message message-log)})))
+                               :raw-body               body})))
             (try
-              (with-open [writer (io/writer temp-file)]
-                (.write writer ^String output))
+              (with-open [in (python-runner/open-output @shared-storage-ref)]
+                (io/copy in temp-file))
               (let [file-size (.length temp-file)]
                 (transforms.instrumentation/with-stage-timing [run-id :file-to-dwh]
                   (transfer-file-to-db driver db transform output-manifest temp-file))

enterprise/backend/src/metabase_enterprise/transforms_python/python_runner.clj

@@ -2,7 +2,6 @@
   (:require
    [clj-http.client :as http]
    [clojure.java.io :as io]
-   [clojure.string :as str]
    [medley.core :as m]
    [metabase-enterprise.transforms-python.s3 :as s3]
    [metabase-enterprise.transforms-python.settings :as transforms-python.settings]
@@ -19,7 +18,7 @@
    [toucan2.core :as t2])
   (:import
    (clojure.lang PersistentQueue)
-   (java.io BufferedWriter File OutputStream OutputStreamWriter)
+   (java.io BufferedWriter File InputStream OutputStream OutputStreamWriter)
    (java.nio.charset StandardCharsets)))
 
 (set! *warn-on-reflection* true)
@@ -203,17 +202,24 @@
                                    {:error string-if-error}))
                                string-if-error)))))
 
-;; temporary, we should not need to realize data/events files into memory longer term
-(defn read-output-objects
-  "Temporary function that strings/jsons stuff in S3 and returns it for compatibility."
+(defn read-events
+  "Returns a vector of event contents (or nil if the event file does not exist)"
   [{:keys [s3-client bucket-name objects]}]
-  (let [{:keys [output output-manifest events]} objects
-        output-content          (s3/read-to-string s3-client bucket-name (:path output) nil)
-        output-manifest-content (s3/read-to-string s3-client bucket-name (:path output-manifest) "{}")
-        events-content          (s3/read-to-string s3-client bucket-name (:path events))]
-    {:output          output-content
-     :output-manifest (json/decode+kw output-manifest-content)
-     :events          (mapv json/decode+kw (str/split-lines events-content))}))
+  ;; note we expect :events to be a small file, limits ought to be enforced by the runner
+  (when-some [in (s3/open-object s3-client bucket-name (:path (:events objects)))]
+    (with-open [in in
+                rdr (io/reader in)]
+      (mapv json/decode+kw (line-seq rdr)))))
+
+(defn read-output-manifest
+  "Return the output manifest map. Returns nil if it does not exist."
+  [{:keys [s3-client bucket-name objects]}]
+  (json/decode+kw (s3/read-to-string s3-client bucket-name (:path (:output-manifest objects)) "{}")))
+
+(defn open-output
+  "Return an InputStream with the output jsonl contents. Close with .close. Returns nil if the object does not exist."
+  [{:keys [s3-client bucket-name objects]}]
+  (s3/open-object s3-client bucket-name (:path (:output objects))))
 
 (defn cancel-python-code-http-call!
   "Calls the /cancel endpoint of the python runner. Returns immediately."

enterprise/backend/src/metabase_enterprise/transforms_python/s3.clj

@@ -4,7 +4,7 @@
    [metabase.util.log :as log])
   (:import
    (clojure.lang IDeref)
-   (java.io Closeable File)
+   (java.io Closeable File ByteArrayInputStream InputStream)
    (java.net URI)
    (java.time Duration)
    (software.amazon.awssdk.auth.credentials AwsBasicCredentials StaticCredentialsProvider)
@@ -183,19 +183,22 @@
   (let [^PutObjectRequest request (put-object-request bucket-name key)]
     (.putObject s3-client request (RequestBody/fromFile file))))
 
-;; TODO optimize our ingestion to stream the data. reading a string will not work in any case for binary files.
+(defn open-object
+  "Get back the contents of the given key as a InputStream."
+  ^InputStream [^S3Client s3-client ^String bucket-name ^String key]
+  (try
+    (let [^GetObjectRequest request (get-object-request bucket-name key)]
+      (.getObject s3-client request))
+    (catch NoSuchKeyException _ nil)))
+
 (defn read-to-string
   "Get back the contents of the given key as a string."
   ([s3-client bucket-name key] (read-to-string s3-client bucket-name key ::throw))
   ([^S3Client s3-client ^String bucket-name ^String key not-found]
-   (try
-     (let [^GetObjectRequest request (get-object-request bucket-name key)
-           response                  (.getObject s3-client request)]
-       (slurp response))
-     (catch NoSuchKeyException e
-       (if (identical? ::throw not-found)
-         (throw e)
-         not-found)))))
+   (if-some [in (open-object s3-client bucket-name key)]
+     (with-open [ret in]
+       (slurp ret))
+     not-found)))
 
 (defn delete
   ;; TODO better error handling

enterprise/backend/test/metabase_enterprise/transforms_python/drivers_test.clj

@@ -29,21 +29,22 @@
   "Execute a Python transform with the given code and tables"
   [{:keys [code tables]}]
   (with-open [shared-storage-ref (s3/open-shared-storage! (or tables {}))]
-    (let [server-url (transforms-python.settings/python-runner-url)
-          cancel-chan (a/promise-chan)
-          table-name->id (or tables {})
-          _ (python-runner/copy-tables-to-s3! {:run-id test-id
-                                               :shared-storage @shared-storage-ref
-                                               :table-name->id table-name->id
-                                               :cancel-chan cancel-chan})
-          response (python-runner/execute-python-code-http-call! {:server-url server-url
-                                                                  :code code
-                                                                  :run-id test-id
-                                                                  :table-name->id table-name->id
-                                                                  :shared-storage @shared-storage-ref})
-          {:keys [output output-manifest events]} (python-runner/read-output-objects @shared-storage-ref)]
+    (let [server-url      (transforms-python.settings/python-runner-url)
+          cancel-chan     (a/promise-chan)
+          table-name->id  (or tables {})
+          _               (python-runner/copy-tables-to-s3! {:run-id         test-id
+                                                             :shared-storage @shared-storage-ref
+                                                             :table-name->id table-name->id
+                                                             :cancel-chan    cancel-chan})
+          response        (python-runner/execute-python-code-http-call! {:server-url     server-url
+                                                                         :code           code
+                                                                         :run-id         test-id
+                                                                         :table-name->id table-name->id
+                                                                         :shared-storage @shared-storage-ref})
+          events          (python-runner/read-events @shared-storage-ref)
+          output-manifest (python-runner/read-output-manifest @shared-storage-ref)]
       (merge (:body response)
-             {:output output
+             {:output (when-some [in (python-runner/open-output @shared-storage-ref)] (with-open [in in] (slurp in)))
               :output-manifest output-manifest
               :stdout (->> events (filter #(= "stdout" (:stream %))) (map :message) (str/join "\n"))
               :stderr (->> events (filter #(= "stderr" (:stream %))) (map :message) (str/join "\n"))}))))

enterprise/backend/test/metabase_enterprise/transforms_python/python_runner_test.clj

@@ -69,10 +69,11 @@
                                                                         :run-id         test-id
                                                                         :table-name->id table-name->id
                                                                         :shared-storage @shared-storage-ref})
-          {:keys [output output-manifest events]} (python-runner/read-output-objects @shared-storage-ref)]
+          events (python-runner/read-events @shared-storage-ref)
+          output-manifest (python-runner/read-output-manifest @shared-storage-ref)]
       ;; not sure about munging this all together but its what tests expect for now
       (merge (:body response)
-             {:output          output
+             {:output          (when-some [in (python-runner/open-output @shared-storage-ref)] (with-open [in in] (slurp in)))
               :output-manifest output-manifest
               :stdout          (->> events (filter #(= "stdout" (:stream %))) (map :message) (str/join "\n"))
               :stderr          (->> events (filter #(= "stderr" (:stream %))) (map :message) (str/join "\n"))}))))

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

@@ -695,8 +695,8 @@
             documents (map #(into {} (map vector col-names %))
                            data)]
         (if (> (bounded-count 2 documents) 1)
-          ;; TODO: batch to avoid realizing everything in memory
-          (mongo.util/insert-many collection documents)
+          (doseq [chunk (partition-all (or driver/*insert-chunk-rows* 1000) documents)]
+            (mongo.util/insert-many collection chunk))
           (mongo.util/insert-one collection (first documents)))))))
 
 ;; Following code is using monger. Leaving it here for a reference as it could be transformed when there is need