feat: PUT when using in-process patching

resolves #35

Author: tekumara

fakesnow/cursor.py

@@ -10,6 +10,7 @@
 from typing import TYPE_CHECKING, Any, cast
 
 import duckdb
+import pandas as pd
 import pyarrow  # needed by fetch_arrow_table()
 import snowflake.connector.converter
 import snowflake.connector.errors
@@ -175,6 +176,9 @@ def execute(
                 transformed = self._transform(exp, params)
                 self._execute(transformed, params)
 
+            if not kwargs.get("server") and (put_stage_data := transformed.args.get("put_stage_data")):  # pyright: ignore[reportPossiblyUnboundVariable]
+                self._put_files(put_stage_data)
+
             return self
         except snowflake.connector.errors.ProgrammingError as e:
             self._sqlstate = e.sqlstate
@@ -188,6 +192,13 @@ def execute(
             msg = f"{e} not implemented. Please raise an issue via https://github.com/tekumara/fakesnow/issues/new"
             raise snowflake.connector.errors.ProgrammingError(msg=msg, errno=9999, sqlstate="99999") from e
 
+    def _put_files(self, put_stage_data: stage.UploadCommandDict) -> None:
+        results = stage.upload_files(put_stage_data)
+        _df = pd.DataFrame.from_records(results)
+        self._duck_conn.execute("select * from _df")
+        self._arrow_table = self._duck_conn.fetch_arrow_table()
+        self._rowcount = self._arrow_table.num_rows
+
     def check_db_and_schema(self, expression: exp.Expression) -> None:
         no_database, no_schema = checks.is_unqualified_table_expression(expression)
 

fakesnow/server.py

@@ -94,7 +94,7 @@ async def query_request(request: Request) -> JSONResponse:
 
         try:
             # only a single sql statement is sent at a time by the python snowflake connector
-            cur = await run_in_threadpool(conn.cursor().execute, sql_text, binding_params=params)
+            cur = await run_in_threadpool(conn.cursor().execute, sql_text, binding_params=params, server=True)
             rowtype = describe_as_rowtype(cur._describe_last_sql())  # noqa: SLF001
 
             expr = cur._last_transformed  # noqa: SLF001

fakesnow/transforms/stage.py

@@ -1,13 +1,16 @@
 from __future__ import annotations
 
 import datetime
+import os
 import tempfile
 from pathlib import PurePath
+from typing import Any, TypedDict
 from urllib.parse import urlparse
 from urllib.request import url2pathname
 
 import snowflake.connector.errors
 import sqlglot
+from snowflake.connector.file_util import SnowflakeFileUtil
 from sqlglot import exp
 
 from fakesnow.params import MutableParams
@@ -16,6 +19,22 @@
 LOCAL_BUCKET_PATH = tempfile.mkdtemp(prefix="fakesnow_bucket_")
 
 
+class StageInfoDict(TypedDict):
+    locationType: str
+    location: str
+    creds: dict[str, Any]
+
+
+class UploadCommandDict(TypedDict):
+    stageInfo: StageInfoDict
+    src_locations: list[str]
+    parallel: int
+    autoCompress: bool
+    sourceCompression: str
+    overwrite: bool
+    command: str
+
+
 def create_stage(
     expression: exp.Expression,
     current_database: str | None,
@@ -240,3 +259,35 @@ def list_stage_files_sql(stage_name: str) -> str:
             strftime(last_modified, '%a, %d %b %Y %H:%M:%S GMT') as last_modified
         from read_blob('{sdir}/*')
     """
+
+
+def upload_files(put_stage_data: UploadCommandDict) -> list[dict[str, Any]]:
+    results = []
+    for src in put_stage_data["src_locations"]:
+        basename = os.path.basename(src)
+        stage_dir = put_stage_data["stageInfo"]["location"]
+
+        os.makedirs(stage_dir, exist_ok=True)
+        gzip_file_name, target_size = SnowflakeFileUtil.compress_file_with_gzip(src, stage_dir)
+
+        # Rename to match expected .gz extension on upload
+        target_basename = basename + ".gz"
+        target = os.path.join(stage_dir, target_basename)
+        os.replace(gzip_file_name, target)
+
+        target_size = os.path.getsize(target)
+        source_size = os.path.getsize(src)
+
+        results.append(
+            {
+                "source": basename,
+                "target": target_basename,
+                "source_size": source_size,
+                "target_size": target_size,
+                "source_compression": "NONE",
+                "target_compression": "GZIP",
+                "status": "UPLOADED",
+                "message": "",
+            }
+        )
+    return results

tests/test_stage.py

@@ -1,8 +1,10 @@
+import os
+import tempfile
 from datetime import timezone
 
 import pytest
 import snowflake.connector.cursor
-from dirty_equals import IsNow
+from dirty_equals import IsDatetime, IsNow
 
 
 def test_create_stage(dcur: snowflake.connector.cursor.SnowflakeCursor):
@@ -111,3 +113,44 @@ def test_create_stage_qmark_quoted(_fakesnow: None):
     ):
         dcur.execute("CREATE STAGE identifier(?)", ('"stage1"',))
         assert dcur.fetchall() == [{"status": "Stage area stage1 successfully created."}]
+
+
+def test_put_list(dcur: snowflake.connector.cursor.DictCursor) -> None:
+    with tempfile.NamedTemporaryFile(mode="w+", suffix=".csv") as temp_file:
+        data = "1,2\n"
+        temp_file.write(data)
+        temp_file.flush()
+        temp_file_path = temp_file.name
+        temp_file_basename = os.path.basename(temp_file_path)
+
+        dcur.execute("CREATE STAGE stage4")
+        dcur.execute(f"PUT 'file://{temp_file_path}' @stage4")
+        assert dcur.fetchall() == [
+            {
+                "source": temp_file_basename,
+                "target": f"{temp_file_basename}.gz",
+                "source_size": len(data),
+                "target_size": 42,  # GZIP compressed size
+                "source_compression": "NONE",
+                "target_compression": "GZIP",
+                "status": "UPLOADED",
+                "message": "",
+            }
+        ]
+
+        dcur.execute("LIST @stage4")
+        results = dcur.fetchall()
+        assert len(results) == 1
+        assert results[0] == {
+            "name": f"stage4/{temp_file_basename}.gz",
+            "size": 42,
+            "md5": "29498d110c32a756df8109e70d22fa36",
+            "last_modified": IsDatetime(
+                # string in RFC 7231 date format (e.g. 'Sat, 31 May 2025 08:50:51 GMT')
+                format_string="%a, %d %b %Y %H:%M:%S GMT"
+            ),
+        }
+
+        # fully qualified stage name quoted
+        dcur.execute('CREATE STAGE db1.schema1."stage5"')
+        dcur.execute(f"PUT 'file://{temp_file_path}' @db1.schema1.\"stage5\"")