fix: transfer PUT files over http in server mode

fakesnow/server.py

@@ -3,17 +3,19 @@
 import gzip
 import json
 import logging
+import os
 import secrets
 from base64 import b64encode
 from dataclasses import dataclass
 from typing import Any
+from urllib.parse import quote
 
 import snowflake.connector.errors
 from sqlglot import Expr, exp, parse_one
 from starlette.applications import Starlette
 from starlette.concurrency import run_in_threadpool
 from starlette.requests import Request
-from starlette.responses import JSONResponse
+from starlette.responses import JSONResponse, Response
 from starlette.routing import Route
 
 from fakesnow import statement_type
@@ -25,7 +27,7 @@
 from fakesnow.instance import FakeSnow
 from fakesnow.rowtype import ColumnInfo, describe_as_rowtype
 from fakesnow.statement_type import DML_TYPE_IDS, statement_type_id
-from fakesnow.transforms import SERVER_VERSION
+from fakesnow.transforms import SERVER_VERSION, stage
 
 logger = logging.getLogger("fakesnow.server")
 # use same format as uvicorn
@@ -149,6 +151,15 @@ async def query_request(request: Request) -> JSONResponse:
             expr = cur._last_transformed  # noqa: SLF001
             assert expr
             if put_stage_data := expr.args.get("put_stage_data"):
+                if not params:
+                    # rewrite the stage info so the client uploads the file over http to this
+                    # server, rather than writing to a local path, which fails when the client
+                    # doesn't share the server's filesystem (eg: the server runs in a container).
+                    # GCS is the only location type the connector uploads to via a plain http
+                    # presigned url, so mimic it. a PUT with a ? placeholder target keeps
+                    # LOCAL_FS because the connector re-requests the presigned url by executing
+                    # the command without bindings (see SnowflakeGCSRestClient).
+                    put_stage_data["stageInfo"] = gcs_stage_info(request, expr, put_stage_data)
                 # this is a PUT command, so return the stage data
                 return JSONResponse(
                     {
@@ -225,6 +236,45 @@ async def query_request(request: Request) -> JSONResponse:
         )
 
 
+def gcs_stage_info(request: Request, expr: Expr, put_stage_data: stage.UploadCommandDict) -> dict[str, Any]:
+    """Stage info that uploads to this server over http via a GCS-style presigned url.
+
+    The connector re-requests the presigned url with the destination file name as the source,
+    so the url targets the file as actually uploaded, even when auto compression renames it.
+    """
+    stage_dir = put_stage_data["stageInfo"]["location"]
+    dst = os.path.basename(put_stage_data["src_locations"][0])
+    path = quote(os.path.relpath(os.path.join(stage_dir, dst), stage.LOCAL_BUCKET_PATH))
+    return {
+        "locationType": "GCS",
+        "location": f"{expr.args['put_stage_name']}/",
+        "creds": {},
+        "presignedUrl": f"{request.url.scheme}://{request.url.netloc}/fs_bucket/{path}",
+    }
+
+
+def _write_bucket_file(path: str, body: bytes) -> str | None:
+    """Write body to path within the local bucket, returning the target path, or None if invalid."""
+    target = os.path.normpath(os.path.join(stage.LOCAL_BUCKET_PATH, path))
+    if not target.startswith(stage.LOCAL_BUCKET_PATH + os.sep):
+        return None
+
+    os.makedirs(os.path.dirname(target), exist_ok=True)
+    with open(target, "wb") as f:
+        f.write(body)
+    return target
+
+
+async def bucket_upload(request: Request) -> Response:
+    """Store an uploaded file in the local bucket, ie: the backing storage for internal stages."""
+    body = await request.body()
+    target = await run_in_threadpool(_write_bucket_file, request.path_params["path"], body)
+    if not target:
+        return Response(status_code=400)
+    logger.info(f"Uploaded {len(body)} bytes to {target}")
+    return Response()
+
+
 def describe_only_response(
     conn: FakeSnowflakeConnection,
     cur: FakeSnowflakeCursor,
@@ -318,6 +368,7 @@ def monitoring_query(request: Request) -> JSONResponse:
         methods=["POST"],
     ),
     Route("/queries/v1/abort-request", lambda _: JSONResponse({"success": True}), methods=["POST"]),
+    Route("/fs_bucket/{path:path}", bucket_upload, methods=["PUT"]),
     Route("/monitoring/queries/{sfqid}", monitoring_query, methods=["GET"]),
 ]
 

fakesnow/transforms/stage.py

@@ -198,7 +198,9 @@ def put_stage(
 
     assert isinstance(expression.this, exp.Literal), "PUT command requires a file path as a literal"
     src_url = urlparse(expression.this.this)
-    src_path = url2pathname(src_url.path)
+    # include netloc to handle relative urls, eg: file://data.csv.gz as sent by the connector
+    # when it re-requests a presigned url using the destination file name
+    src_path = src_url.netloc + url2pathname(src_url.path)
     target = expression.args["target"]
 
     assert isinstance(target, exp.Var), f"{target} is not a exp.Var"

tests/test_server.py

@@ -437,6 +437,68 @@ def test_server_put_non_existent_stage(sdcur: snowflake.connector.cursor.DictCur
         )
 
 
+def test_server_put_presigned_url(server: dict, sconn: snowflake.connector.SnowflakeConnection) -> None:
+    # in server mode the client may not share the server's filesystem, so PUT hands back a
+    # GCS-style presigned url and the client uploads over http to this server, which stores
+    # the file in the stage
+    conn = sconn
+    with (
+        conn.cursor(snowflake.connector.cursor.DictCursor) as dcur,
+        tempfile.NamedTemporaryFile(suffix=".csv") as temp_file,
+    ):
+        temp_file_basename = os.path.basename(temp_file.name)
+        dcur.execute("CREATE STAGE presigned_stage")
+
+        result = conn.cmd_query(
+            f"PUT 'file://{temp_file.name}' @presigned_stage",
+            conn._next_sequence_counter(),  # noqa: SLF001
+            uuid.uuid4(),
+        )
+
+        stage_info = result["data"]["stageInfo"]
+        assert stage_info["locationType"] == "GCS"
+        assert stage_info["presignedUrl"] == (
+            f"http://{server['host']}:{server['port']}/fs_bucket/DB1/SCHEMA1/PRESIGNED_STAGE/{temp_file_basename}"
+        )
+
+        data = b"1,2\n"
+        response = requests.put(stage_info["presignedUrl"], data=data, timeout=5)
+        assert response.status_code == 200
+
+        dcur.execute("LIST @presigned_stage")
+        assert dcur.fetchall() == [
+            {
+                "name": f"presigned_stage/{temp_file_basename}",
+                "size": len(data),
+                "md5": IsStr(regex=r"^[0-9a-f]{32}$"),
+                "last_modified": IsDatetime(format_string="%a, %d %b %Y %H:%M:%S GMT"),
+            }
+        ]
+
+
+def test_server_put_qmark_target_stays_local(sconn: snowflake.connector.SnowflakeConnection) -> None:
+    # the connector re-requests a presigned url by executing the PUT without bindings, which
+    # cannot resolve a ? target, so a bound target keeps the local filesystem upload
+    conn = sconn
+    with conn.cursor() as cur, tempfile.NamedTemporaryFile(suffix=".csv") as temp_file:
+        cur.execute("CREATE STAGE qmark_stage")
+
+        result = conn.cmd_query(
+            f"PUT 'file://{temp_file.name}' ?",
+            conn._next_sequence_counter(),  # noqa: SLF001
+            uuid.uuid4(),
+            binding_params={"1": {"type": "TEXT", "value": "@qmark_stage"}},
+        )
+
+        assert result["data"]["stageInfo"]["locationType"] == "LOCAL_FS"
+
+
+def test_server_bucket_upload_rejects_path_outside_bucket(server: dict) -> None:
+    response = requests.put(f"http://{server['host']}:{server['port']}/fs_bucket//etc/passwd", data=b"x", timeout=5)
+
+    assert response.status_code == 400
+
+
 def test_server_response_params(server: dict) -> None:
     # mimic the jdbc driver
     headers = {