fakesnow 0.9.38__py3-none-any.whl → 0.9.40__py3-none-any.whl

fakesnow/checks.py

@@ -39,8 +39,8 @@ def is_unqualified_table_expression(expression: exp.Expression) -> tuple[bool, b
             # "CREATE/DROP SCHEMA"
             no_database = not node.args.get("catalog")
             no_schema = False
-        elif parent_kind.upper() in {"TABLE", "VIEW"}:
-            # "CREATE/DROP TABLE/VIEW"
+        elif parent_kind.upper() in {"TABLE", "VIEW", "STAGE"}:
+            # "CREATE/DROP TABLE/VIEW/STAGE"
             no_database = not node.args.get("catalog")
             no_schema = not node.args.get("db")
         else:

fakesnow/copy_into.py

@@ -1,8 +1,9 @@
 from __future__ import annotations
 
+import datetime
 from collections.abc import Sequence
 from dataclasses import dataclass, field
-from typing import Any, Protocol, cast
+from typing import Any, NamedTuple, Protocol, cast
 from urllib.parse import urlparse, urlunparse
 
 import duckdb
@@ -10,30 +11,125 @@ import snowflake.connector.errors
 from duckdb import DuckDBPyConnection
 from sqlglot import exp
 
+import fakesnow.transforms.stage as stage
 from fakesnow import logger
 
 
+class LoadHistoryRecord(NamedTuple):
+    """Represents a record in the INFORMATION_SCHEMA.LOAD_HISTORY table."""
+
+    schema_name: str
+    file_name: str
+    table_name: str
+    last_load_time: str  # ISO8601 datetime with timezone
+    status: str
+    row_count: int
+    row_parsed: int
+    first_error_message: str | None
+    first_error_line_number: int | None
+    first_error_character_position: int | None
+    first_error_col_name: str | None
+    error_count: int
+    error_limit: int | None
+
+
 def copy_into(
-    duck_conn: DuckDBPyConnection, expr: exp.Copy, params: Sequence[Any] | dict[Any, Any] | None = None
+    duck_conn: DuckDBPyConnection,
+    current_database: str | None,
+    current_schema: str | None,
+    expr: exp.Copy,
+    params: Sequence[Any] | dict[Any, Any] | None = None,
 ) -> str:
     cparams = _params(expr)
-    urls = _source_urls(expr, cparams.files)
+    if isinstance(cparams.file_format, ReadParquet):
+        from_ = expr.args["files"][0]
+        # parquet must use MATCH_BY_COLUMN_NAME (TODO) or a copy transformation
+        # ie: the from clause in COPY INTO must be a subquery
+        if not isinstance(from_, exp.Subquery):
+            raise snowflake.connector.errors.ProgrammingError(
+                msg="SQL compilation error:\nPARQUET file format can produce one and only one column of type variant, object, or array. Load data into separate columns using the MATCH_BY_COLUMN_NAME copy option or copy with transformation.",  # noqa: E501
+                errno=2019,
+                sqlstate="0A000",
+            )
+
+    from_source = _from_source(expr)
+    source = (
+        stage_url_from_var(from_source, duck_conn, current_database, current_schema)
+        if from_source.startswith("@")
+        else from_source
+    )
+    urls = _source_urls(source, cparams.files)
+
     inserts = _inserts(expr, cparams, urls)
+    table = expr.this
+    if isinstance(expr.this, exp.Table):
+        table = expr.this
+    elif isinstance(expr.this, exp.Schema) and isinstance(expr.this.this, exp.Table):
+        table = expr.this.this
+    else:
+        raise AssertionError(f"copy into {expr.this.__class__} is not Table or Schema")
+
+    schema = table.db or current_schema
+    assert schema
 
-    results = []
+    histories: list[LoadHistoryRecord] = []
+    load_time = datetime.datetime.now(datetime.timezone.utc).isoformat()
     try:
-        # TODO: fetch files last modified dates and check if file exists in load_history already
+        check_sql = "SELECT 1 FROM _fs_information_schema._fs_load_history WHERE FILE_NAME = ? LIMIT 1"
+
         for i, url in zip(inserts, urls):
-            sql = i.sql(dialect="duckdb")
-            logger.log_sql(sql, params)
+            # Check if file has been loaded into any table before
+            duck_conn.execute(check_sql, [url])
+            if duck_conn.fetchone() and not cparams.force:
+                affected_count = 0
+                status = "LOAD_SKIPPED"
+                error_limit = None
+                error_count = 1
+                first_error_message = "File was loaded before."
+            else:
+                sql = i.sql(dialect="duckdb")
+                logger.log_sql(sql, params)
+                duck_conn.execute(sql, params)
+                (affected_count,) = duck_conn.fetchall()[0]
+                status = "LOADED"
+                error_limit = 1
+                error_count = 0
+                first_error_message = None
+
+            history = LoadHistoryRecord(
+                schema_name=schema,
+                file_name=url,
+                table_name=table.name,
+                last_load_time=load_time,
+                status=status,
+                row_count=affected_count,
+                row_parsed=affected_count,
+                first_error_message=first_error_message,
+                first_error_line_number=None,
+                first_error_character_position=None,
+                first_error_col_name=None,
+                error_count=error_count,
+                error_limit=error_limit,
+            )
+            histories.append(history)
+
+        if insert_histories := [h for h in histories if h.status != "LOAD_SKIPPED"]:
+            values = "\n ,".join(str(tuple(history)).replace("None", "NULL") for history in insert_histories)
+            sql = f"INSERT INTO _fs_information_schema._fs_load_history VALUES {values}"
             duck_conn.execute(sql, params)
-            (affected_count,) = duck_conn.fetchall()[0]
-            results.append(f"('{url}', 'LOADED', {affected_count}, {affected_count}, 1, 0, NULL, NULL, NULL, NULL)")
 
-            # TODO: update load_history with the results if loaded
-
-        columns = "file, status, rows_parsed, rows_loaded, error_limit, errors_seen, first_error, first_error_line, first_error_character, first_error_column_name"  # noqa: E501
-        values = "\n, ".join(results)
+        columns = (
+            "file, status, rows_parsed, rows_loaded, error_limit, errors_seen, first_error, first_error_line, "
+            "first_error_character, first_error_column_name"
+        )
+        values = "\n, ".join(
+            f"('{h.file_name}', '{h.status}', {h.row_parsed}, {h.row_count}, "
+            f"{h.error_limit or 'NULL'}, {h.error_count}, "
+            f"{repr(h.first_error_message) if h.first_error_message else 'NULL'}, "
+            f"{h.first_error_line_number or 'NULL'}, {h.first_error_character_position or 'NULL'}, "
+            f"{h.first_error_col_name or 'NULL'})"
+            for h in histories
+        )
         sql = f"SELECT * FROM (VALUES\n  {values}\n) AS t({columns})"
         duck_conn.execute(sql)
         return sql
@@ -62,6 +158,8 @@ def _params(expr: exp.Copy) -> Params:
 
             if var_type == "CSV":
                 kwargs["file_format"] = handle_csv(param.expressions)
+            elif var_type == "PARQUET":
+                kwargs["file_format"] = ReadParquet()
             else:
                 raise NotImplementedError(f"{var_type} FILE_FORMAT is not currently implemented")
         elif var == "FORCE":
@@ -71,28 +169,67 @@ def _params(expr: exp.Copy) -> Params:
         else:
             raise ValueError(f"Unknown copy parameter: {param.this}")
 
-    if not force:
-        raise NotImplementedError("COPY INTO with FORCE=false (default) is not currently implemented")
+    return Params(force=force, **kwargs)
 
-    return Params(**kwargs)
 
+def _from_source(expr: exp.Copy) -> str:
+    # NB: sqlglot parses the from clause as "files" strangely
+    from_ = expr.args["files"][0].this
+
+    if isinstance(from_, exp.Select):
+        from_table = from_.args["from"]
+        # if a subquery is used in the FROM clause it must be loaded from a stage not an external location
+        assert isinstance(from_table, exp.From), f"{from_table.__class__} is not a From"
+        assert isinstance(from_table.this, exp.Table), f"{from_table.__class__} is not a Table"
+        var = from_table.this.this
+        if not isinstance(var, exp.Var):
+            # not a very helpful message, but this is what Snowflake returns
+            raise snowflake.connector.errors.ProgrammingError(
+                msg=f"SQL compilation error:\ninvalid URL prefix found in: {from_table.this.this}",
+                errno=1011,
+                sqlstate="42601",
+            )
+        # return the name of the stage, eg: @stage1
+        return var.this
+
+    assert isinstance(from_, exp.Literal), f"{from_} is not a exp.Literal"
+    # return url
+    return from_.name
+
+
+def stage_url_from_var(
+    from_source: str, duck_conn: DuckDBPyConnection, current_database: str | None, current_schema: str | None
+) -> str:
+    database_name, schema_name, name = stage.parts_from_var(from_source, current_database, current_schema)
+
+    # Look up the stage URL
+    duck_conn.execute(
+        """
+        SELECT url FROM _fs_global._fs_information_schema._fs_stages
+        WHERE database_name = ? and schema_name  = ? and name = ?
+        """,
+        (database_name, schema_name, name),
+    )
+    if result := duck_conn.fetchone():
+        return result[0]
+    else:
+        raise snowflake.connector.errors.ProgrammingError(
+            msg=f"SQL compilation error:\nStage '{database_name}.{schema_name}.{name}' does not exist or not authorized.",  # noqa: E501
+            errno=2003,
+            sqlstate="02000",
+        )
 
-def _source_urls(expr: exp.Copy, files: list[str]) -> list[str]:
-    """
-    Given a COPY statement and a list of files, return a list of URLs with each file appended as a fragment.
-    Checks that the source is a valid URL.
-    """
-    source = expr.args["files"][0].this
-    assert isinstance(source, exp.Literal), f"{source} is not a exp.Literal"
 
-    scheme, netloc, path, params, query, fragment = urlparse(source.name)
+def _source_urls(from_source: str, files: list[str]) -> list[str]:
+    """Convert from_source to a list of URLs."""
+    scheme, netloc, path, params, query, fragment = urlparse(from_source)
     if not scheme:
         raise snowflake.connector.errors.ProgrammingError(
-            msg=f"SQL compilation error:\ninvalid URL prefix found in: '{source.name}'", errno=1011, sqlstate="42601"
+            msg=f"SQL compilation error:\ninvalid URL prefix found in: '{from_source}'", errno=1011, sqlstate="42601"
         )
 
     # rebuild url from components to ensure correct handling of host slash
-    return [_urlunparse(scheme, netloc, path, params, query, fragment, file) for file in files] or [source.name]
+    return [_urlunparse(scheme, netloc, path, params, query, fragment, file) for file in files] or [from_source]
 
 
 def _urlunparse(scheme: str, netloc: str, path: str, params: str, query: str, fragment: str, suffix: str) -> str:
@@ -111,9 +248,16 @@ def _urlunparse(scheme: str, netloc: str, path: str, params: str, query: str, fr
 def _inserts(expr: exp.Copy, params: Params, urls: list[str]) -> list[exp.Expression]:
     # INTO expression
     target = expr.this
-    columns = [exp.Column(this=exp.Identifier(this=f"column{i}")) for i in range(len(target.expressions))] or [
-        exp.Column(this=exp.Star())
-    ]
+
+    from_ = expr.args["files"][0]
+    if isinstance(from_, exp.Subquery):
+        select = from_.this
+        assert isinstance(select, exp.Select), f"{select.__class__} is not a Select"
+        columns = _strip_json_extract(select).expressions
+    else:
+        columns = [exp.Column(this=exp.Identifier(this=f"column{i}")) for i in range(len(target.expressions))] or [
+            exp.Column(this=exp.Star())
+        ]
 
     return [
         exp.Insert(
@@ -124,6 +268,20 @@ def _inserts(expr: exp.Copy, params: Params, urls: list[str]) -> list[exp.Expres
     ]
 
 
+def _strip_json_extract(expr: exp.Select) -> exp.Select:
+    """
+    Strip $1 prefix from SELECT statement columns.
+    """
+    dollar1 = exp.Parameter(this=exp.Literal(this="1", is_string=False))
+
+    for p in expr.find_all(exp.Parameter):
+        if p == dollar1 and p.parent and (key := p.parent.expression.find(exp.JSONPathKey)):
+            assert p.parent.parent, expr
+            p.parent.parent.args["this"] = exp.Identifier(this=key.this)
+
+    return expr
+
+
 def handle_csv(expressions: list[exp.Property]) -> ReadCSV:
     skip_header = ReadCSV.skip_header
     quote = ReadCSV.quote
@@ -192,8 +350,15 @@ class ReadCSV(FileTypeHandler):
         return exp.func("read_csv", exp.Literal(this=url, is_string=True), *args)
 
 
+@dataclass
+class ReadParquet(FileTypeHandler):
+    def read_expression(self, url: str) -> exp.Expression:
+        return exp.func("read_parquet", exp.Literal(this=url, is_string=True))
+
+
 @dataclass
 class Params:
     files: list[str] = field(default_factory=list)
     # Snowflake defaults to CSV when no file format is specified
     file_format: FileTypeHandler = field(default_factory=ReadCSV)
+    force: bool = False

fakesnow/cursor.py

@@ -43,6 +43,7 @@ SQL_CREATED_DATABASE = Template("SELECT 'Database ${name} successfully created.'
 SQL_CREATED_SCHEMA = Template("SELECT 'Schema ${name} successfully created.' as 'status'")
 SQL_CREATED_TABLE = Template("SELECT 'Table ${name} successfully created.' as 'status'")
 SQL_CREATED_VIEW = Template("SELECT 'View ${name} successfully created.' as 'status'")
+SQL_CREATED_STAGE = Template("SELECT 'Stage area ${name} successfully created.' as status")
 SQL_DROPPED = Template("SELECT '${name} successfully dropped.' as 'status'")
 SQL_INSERTED_ROWS = Template("SELECT ${count} as 'number of rows inserted'")
 SQL_UPDATED_ROWS = Template("SELECT ${count} as 'number of rows updated', 0 as 'number of multi-joined rows updated'")
@@ -75,6 +76,7 @@ class FakeSnowflakeCursor:
         self._use_dict_result = use_dict_result
         self._last_sql = None
         self._last_params = None
+        self._last_transformed = None
         self._sqlstate = None
         self._arraysize = 1
         self._arrow_table = None
@@ -105,6 +107,7 @@ class FakeSnowflakeCursor:
     def close(self) -> bool:
         self._last_sql = None
         self._last_params = None
+        self._last_transformed = None
         return True
 
     def describe(self, command: str, *args: Any, **kwargs: Any) -> list[ResultMetadata]:
@@ -203,10 +206,8 @@ class FakeSnowflakeCursor:
             .transform(transforms.create_database, db_path=self._conn.db_path)
             .transform(transforms.extract_comment_on_table)
             .transform(transforms.extract_comment_on_columns)
-            .transform(transforms.information_schema_fs_columns)
+            .transform(transforms.information_schema_fs)
             .transform(transforms.information_schema_databases, current_schema=self._conn.schema)
-            .transform(transforms.information_schema_fs_tables)
-            .transform(transforms.information_schema_fs_views)
             .transform(transforms.drop_schema_cascade)
             .transform(transforms.tag)
             .transform(transforms.semi_structured_types)
@@ -240,6 +241,7 @@ class FakeSnowflakeCursor:
             .transform(transforms.identifier)
             .transform(transforms.array_agg_within_group)
             .transform(transforms.array_agg)
+            .transform(transforms.array_construct_etc)
             .transform(transforms.dateadd_date_cast)
             .transform(transforms.dateadd_string_literal_timestamp_cast)
             .transform(transforms.datediff_string_literal_timestamp_cast)
@@ -248,8 +250,9 @@ class FakeSnowflakeCursor:
             .transform(transforms.show_procedures)
             .transform(transforms.show_warehouses)
             .transform(lambda e: transforms.show_schemas(e, self._conn.database))
-            .transform(lambda e: transforms.show_tables_etc(e, self._conn.database))
-            .transform(lambda e: transforms.show_columns(e, self._conn.database))
+            .transform(lambda e: transforms.show_tables_etc(e, self._conn.database, self._conn.schema))
+            .transform(lambda e: transforms.show_columns(e, self._conn.database, self._conn.schema))
+            .transform(lambda e: transforms.show_stages(e, self._conn.database, self._conn.schema))
             # TODO collapse into a single show_keys function
             .transform(lambda e: transforms.show_keys(e, self._conn.database, kind="PRIMARY"))
             .transform(lambda e: transforms.show_keys(e, self._conn.database, kind="UNIQUE"))
@@ -260,6 +263,8 @@ class FakeSnowflakeCursor:
             .transform(transforms.create_clone)
             .transform(transforms.alias_in_join)
             .transform(transforms.alter_table_strip_cluster_by)
+            .transform(lambda e: transforms.create_stage(e, self._conn.database, self._conn.schema))
+            .transform(lambda e: transforms.put_stage(e, self._conn.database, self._conn.schema))
         )
 
     def _transform_explode(self, expression: exp.Expression) -> list[exp.Expression]:
@@ -287,17 +292,17 @@ class FakeSnowflakeCursor:
 
         try:
             if isinstance(transformed, exp.Copy):
-                sql = copy_into(self._duck_conn, transformed, params)
+                sql = copy_into(self._duck_conn, self._conn.database, self._conn.schema, transformed, params)
             else:
                 logger.log_sql(sql, params)
                 self._duck_conn.execute(sql, params)
         except duckdb.BinderException as e:
             msg = e.args[0]
-            raise snowflake.connector.errors.ProgrammingError(msg=msg, errno=2043, sqlstate="02000") from None
+            raise snowflake.connector.errors.ProgrammingError(msg=msg, errno=2043, sqlstate="02000") from e
         except duckdb.CatalogException as e:
             # minimal processing to make it look like a snowflake exception, message content may differ
             msg = cast(str, e.args[0]).split("\n")[0]
-            raise snowflake.connector.errors.ProgrammingError(msg=msg, errno=2003, sqlstate="42S02") from None
+            raise snowflake.connector.errors.ProgrammingError(msg=msg, errno=2003, sqlstate="42S02") from e
         except duckdb.TransactionException as e:
             if "cannot rollback - no transaction is active" in str(
                 e
@@ -307,9 +312,9 @@ class FakeSnowflakeCursor:
             else:
                 raise e
         except duckdb.ConnectionException as e:
-            raise snowflake.connector.errors.DatabaseError(msg=e.args[0], errno=250002, sqlstate="08003") from None
+            raise snowflake.connector.errors.DatabaseError(msg=e.args[0], errno=250002, sqlstate="08003") from e
         except duckdb.ParserException as e:
-            raise snowflake.connector.errors.ProgrammingError(msg=e.args[0], errno=1003, sqlstate="42000") from None
+            raise snowflake.connector.errors.ProgrammingError(msg=e.args[0], errno=1003, sqlstate="42000") from e
 
         affected_count = None
 
@@ -329,6 +334,15 @@ class FakeSnowflakeCursor:
             self._duck_conn.execute(info_schema.per_db_creation_sql(create_db_name))
             result_sql = SQL_CREATED_DATABASE.substitute(name=create_db_name)
 
+        elif stage_name := transformed.args.get("create_stage_name"):
+            if stage_name == "?":
+                assert isinstance(params, (tuple, list)) and len(params) == 1, (
+                    "Expected single parameter for create_stage_name"
+                )
+                result_sql = SQL_CREATED_STAGE.substitute(name=params[0].upper())
+            else:
+                result_sql = SQL_CREATED_STAGE.substitute(name=stage_name.upper())
+
         elif cmd == "INSERT":
             (affected_count,) = self._duck_conn.fetchall()[0]
             result_sql = SQL_INSERTED_ROWS.substitute(count=affected_count)
@@ -401,8 +415,21 @@ class FakeSnowflakeCursor:
         self._rowcount = affected_count or self._arrow_table.num_rows
         self._sfqid = str(uuid.uuid4())
 
+        if stage_name := transformed.args.get("put_stage_name"):
+            if stage_name == "?":
+                assert isinstance(params, (tuple, list)) and len(params) == 1, (
+                    "Expected single parameter for put_stage_name"
+                )
+            if self._arrow_table.num_rows != 1:
+                raise snowflake.connector.errors.ProgrammingError(
+                    msg=f"SQL compilation error:\nStage '{stage_name}' does not exist or not authorized.",
+                    errno=2003,
+                    sqlstate="02000",
+                )
+
         self._last_sql = result_sql or sql
         self._last_params = None if result_sql else params
+        self._last_transformed = transformed
 
     def executemany(
         self,

fakesnow/info_schema.py

@@ -185,6 +185,49 @@ where database_name = '${catalog}'
 """
 )
 
+SQL_CREATE_LOAD_HISTORY_TABLE = Template(
+    """
+create table if not exists ${catalog}._fs_information_schema._fs_load_history (
+    SCHEMA_NAME VARCHAR,
+    FILE_NAME VARCHAR,
+    TABLE_NAME VARCHAR,
+    LAST_LOAD_TIME TIMESTAMPTZ,
+    STATUS VARCHAR,
+    ROW_COUNT INTEGER,
+    ROW_PARSED INTEGER,
+    FIRST_ERROR_MESSAGE VARCHAR,
+    FIRST_ERROR_LINE_NUMBER INTEGER,
+    FIRST_ERROR_CHARACTER_POSITION INTEGER,
+    FIRST_ERROR_COL_NAME VARCHAR,
+    ERROR_COUNT INTEGER,
+    ERROR_LIMIT INTEGER
+)
+    """
+)
+
+
+SQL_CREATE_GLOBAL_INFORMATION_SCHEMA_STAGES_TABLE = """
+CREATE TABLE IF NOT EXISTS _fs_global._fs_information_schema._fs_stages (
+    created_on TIMESTAMPTZ,
+    name TEXT,
+    database_name TEXT,
+    schema_name TEXT,
+    url TEXT,
+    has_credentials TEXT,
+    has_encryption_key TEXT,
+    owner TEXT,
+    comment TEXT,
+    region TEXT,
+    type TEXT,
+    cloud TEXT,
+    notification_channel TEXT,
+    storage_integration TEXT,
+    endpoint TEXT,
+    owner_role_type TEXT,
+    directory_enabled TEXT
+);
+"""
+
 
 def per_db_creation_sql(catalog: str) -> str:
     return f"""
@@ -193,6 +236,7 @@ def per_db_creation_sql(catalog: str) -> str:
         {SQL_CREATE_INFORMATION_SCHEMA_DATABASES_VIEW.substitute(catalog=catalog)};
         {SQL_CREATE_INFORMATION_SCHEMA_TABLES_VIEW.substitute(catalog=catalog)};
         {SQL_CREATE_INFORMATION_SCHEMA_VIEWS_VIEW.substitute(catalog=catalog)};
+        {SQL_CREATE_LOAD_HISTORY_TABLE.substitute(catalog=catalog)};
     """
 
 
@@ -203,6 +247,7 @@ def fs_global_creation_sql() -> str:
         {SQL_CREATE_GLOBAL_INFORMATION_SCHEMA_COLUMNS_EXT};
         {SQL_CREATE_GLOBAL_INFORMATION_SCHEMA_COLUMNS_VIEW};
         {SQL_CREATE_GLOBAL_INFORMATION_SCHEMA_USERS_TABLE};
+        {SQL_CREATE_GLOBAL_INFORMATION_SCHEMA_STAGES_TABLE}
     """
 
 

fakesnow/macros.py

@@ -29,9 +29,20 @@ CREATE OR REPLACE MACRO ${catalog}._fs_flatten(input) AS TABLE
     """
 )
 
+# emulates https://docs.snowflake.com/en/sql-reference/functions/array_construct_compact
+# requires transforms.array_construct_compact
+ARRAY_CONSTRUCT_COMPACT = Template(
+    """
+CREATE OR REPLACE MACRO ${catalog}.array_construct_compact(list) AS (
+    SELECT ARRAY_AGG(x)::JSON FROM UNNEST(list) AS t(x) WHERE x IS NOT NULL
+);
+"""
+)
+
 
 def creation_sql(catalog: str) -> str:
     return f"""
         {EQUAL_NULL.substitute(catalog=catalog)};
         {FS_FLATTEN.substitute(catalog=catalog)};
+        {ARRAY_CONSTRUCT_COMPACT.substitute(catalog=catalog)};
     """

fakesnow/server.py

@@ -9,6 +9,7 @@ from dataclasses import dataclass
 from typing import Any
 
 import snowflake.connector.errors
+from sqlglot import parse_one
 from starlette.applications import Starlette
 from starlette.concurrency import run_in_threadpool
 from starlette.requests import Request
@@ -63,6 +64,7 @@ async def login_request(request: Request) -> JSONResponse:
                     {"name": "AUTOCOMMIT", "value": True},
                     {"name": "CLIENT_SESSION_KEEP_ALIVE_HEARTBEAT_FREQUENCY", "value": 3600},
                 ],
+                "sessionInfo": {},
             },
             "success": True,
         }
@@ -88,11 +90,24 @@ async def query_request(request: Request) -> JSONResponse:
         else:
             params = None
 
+        expr = parse_one(sql_text, read="snowflake")
+
         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)
             rowtype = describe_as_rowtype(cur._describe_last_sql())  # noqa: SLF001
 
+            expr = cur._last_transformed  # noqa: SLF001
+            assert expr
+            if put_stage_data := expr.args.get("put_stage_data"):
+                # this is a PUT command, so return the stage data
+                return JSONResponse(
+                    {
+                        "data": put_stage_data,
+                        "success": True,
+                    }
+                )
+
         except snowflake.connector.errors.ProgrammingError as e:
             logger.info(f"{sql_text=} ProgrammingError {e}")
             code = f"{e.errno:06d}"

fakesnow/transforms/__init__.py

@@ -8,16 +8,22 @@ from fakesnow.transforms.show import (
     show_keys as show_keys,
     show_procedures as show_procedures,
     show_schemas as show_schemas,
+    show_stages as show_stages,
     show_tables_etc as show_tables_etc,
     show_users as show_users,
     show_warehouses as show_warehouses,
 )
+from fakesnow.transforms.stage import (
+    create_stage as create_stage,
+    put_stage as put_stage,
+)
 from fakesnow.transforms.transforms import (
     SUCCESS_NOP as SUCCESS_NOP,
     alias_in_join as alias_in_join,
     alter_table_strip_cluster_by as alter_table_strip_cluster_by,
     array_agg as array_agg,
     array_agg_within_group as array_agg_within_group,
+    array_construct_etc as array_construct_etc,
     array_size as array_size,
     create_clone as create_clone,
     create_database as create_database,
@@ -36,9 +42,7 @@ from fakesnow.transforms.transforms import (
     identifier as identifier,
     indices_to_json_extract as indices_to_json_extract,
     information_schema_databases as information_schema_databases,
-    information_schema_fs_columns as information_schema_fs_columns,
-    information_schema_fs_tables as information_schema_fs_tables,
-    information_schema_fs_views as information_schema_fs_views,
+    information_schema_fs as information_schema_fs,
     integer_precision as integer_precision,
     json_extract_cased_as_varchar as json_extract_cased_as_varchar,
     json_extract_cast_as_varchar as json_extract_cast_as_varchar,