test: test for parquet output (#86)

wheretrue · Jan 25, 2024 · 050cd5c · 050cd5c
1 parent 01a65ab
commit 050cd5c
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 0 deletions.
diff --git a/bin/test.sh b/bin/test.sh
@@ -57,4 +57,10 @@ aws --endpoint-url=http://localhost:4566 s3api put-object --bucket test-bucket -
 # Make the bucket public.
 aws --endpoint-url=http://localhost:4566 s3api put-bucket-acl --bucket test-bucket --acl public-read
 
+# Create the parquet-bucket.
+aws --endpoint-url=http://localhost:4566 s3 mb s3://parquet-bucket
+
+# Make the bucket public.
+aws --endpoint-url=http://localhost:4566 s3api put-bucket-acl --bucket parquet-bucket --acl public-read
+
 pytest
diff --git a/python/tests/test_session.py b/python/tests/test_session.py
@@ -111,3 +111,18 @@ def test_read_from_s3():
     arrow_table = session.sql(query).to_arrow()
 
     assert len(arrow_table) == 2
+
+
+def test_copy_to_s3():
+    """Test copying to s3."""
+    session = connect()
+
+    s3_input_path = "s3://test-bucket/test.fasta"
+    parquet_output = "s3://parquet-bucket/test.parquet"
+
+    query = f"COPY (SELECT * FROM fasta_scan('{s3_input_path}')) TO '{parquet_output}' (FORMAT PARQUET)"
+
+    session.register_object_store_from_url(parquet_output)
+
+    # Should not raise an exception
+    session.execute(query)