Skip to content

Commit

Permalink
feat(snowflake): read_csv with https
Browse files Browse the repository at this point in the history
  • Loading branch information
IndexSeek authored and cpcloud committed Dec 18, 2023
1 parent 233dce1 commit 72752eb
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 4 deletions.
16 changes: 12 additions & 4 deletions ibis/backends/snowflake/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
import warnings
from pathlib import Path
from typing import TYPE_CHECKING, Any
from urllib.request import urlretrieve

import pyarrow as pa
import pyarrow_hotfix # noqa: F401
Expand Down Expand Up @@ -777,10 +778,17 @@ def read_csv(
)
con.exec_driver_sql(create_infer_fmt)

# copy the local file to the stage
con.exec_driver_sql(
f"PUT 'file://{Path(path).absolute()}' @{stage} PARALLEL = {threads:d}"
)
if path.startswith("https://"):
with tempfile.NamedTemporaryFile() as tmp:
urlretrieve(path, filename=tmp.name)
tmp.flush()
con.exec_driver_sql(
f"PUT 'file://{tmp.name}' @{stage} PARALLEL = {threads:d} AUTO_COMPRESS = TRUE"
)
else:
con.exec_driver_sql(
f"PUT 'file://{Path(path).absolute()}' @{stage} PARALLEL = {threads:d} AUTO_COMPRESS = TRUE"
)

# handle setting up the schema in python because snowflake is
# broken for csv globs: it cannot parse the result of the following
Expand Down
16 changes: 16 additions & 0 deletions ibis/backends/snowflake/tests/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -197,6 +197,22 @@ def test_read_csv_options(con, tmp_path):
assert t.schema() == ibis.schema(dict(a="int64", b="int64"))


def test_read_csv_https(con):
t = con.read_csv(
"https://storage.googleapis.com/ibis-tutorial-data/wowah_data/locations.csv",
field_optionally_enclosed_by='"',
)
assert t.schema() == ibis.schema(
{
"Map_ID": "int64",
"Location_Type": "string",
"Location_Name": "string",
"Game_Version": "string",
}
)
assert t.count().execute() == 151


@pytest.fixture(scope="module")
def json_data():
return [
Expand Down

0 comments on commit 72752eb

Please sign in to comment.