def test_put_get_with_gcp(tmpdir, conn_cnx, db_parameters, is_public_test, enable_gcs_downscoped, from_path):
"""[gcp] Puts and Gets a small text using gcp."""
if enable_gcs_downscoped and is_public_test:
pytest.xfail("Server need to update with merged change. Expected release version: 4.41.0")
# create a data file
fname = str(tmpdir.join('test_put_get_with_gcp_token.txt.gz'))
original_contents = "123,test1\n456,test2\n"
with gzip.open(fname, 'wb') as f:
f.write(original_contents.encode(UTF8))
tmp_dir = str(tmpdir.mkdir('test_put_get_with_gcp_token'))
table_name = random_string(5, 'snow32806_')
with conn_cnx() as cnx:
with cnx.cursor() as csr:
try:
csr.execute(f'ALTER SESSION SET GCS_USE_DOWNSCOPED_CREDENTIAL = {enable_gcs_downscoped}')
except ProgrammingError as e:
if enable_gcs_downscoped:
# not raise error when the parameter is not available yet, using old behavior
raise e
csr.execute("create or replace table {} (a int, b string)".format(table_name))
try:
file_stream = None if from_path else open(fname, 'rb')
put(csr, fname, f"%{table_name}", from_path,
sql_options=" auto_compress=true parallel=30",
file_stream=file_stream)
assert csr.fetchone()[6] == 'UPLOADED'
csr.execute("copy into {}".format(table_name))
csr.execute("rm @%{}".format(table_name))
assert csr.execute("ls @%{}".format(table_name)).fetchall() == []
csr.execute("copy into @%{table_name} from {table_name} "
"file_format=(type=csv compression='gzip')".format(table_name=table_name))
csr.execute("get @%{table_name} file://{}".format(tmp_dir, table_name=table_name))
rec = csr.fetchone()
assert rec[0].startswith('data_'), 'A file downloaded by GET'
assert rec[1] == 36, 'Return right file size'
assert rec[2] == 'DOWNLOADED', 'Return DOWNLOADED status'
assert rec[3] == '', 'Return no error message'
finally:
if file_stream:
file_stream.close()
csr.execute("drop table {}".format(table_name))
files = glob.glob(os.path.join(tmp_dir, 'data_*'))
with gzip.open(files[0], 'rb') as fd:
contents = fd.read().decode(UTF8)
assert original_contents == contents, 'Output is different from the original file'
def mocked_put(*args, **kwargs):
if mocked_put.counter == 0:
exc = requests.exceptions.HTTPError(response=requests.Response())
exc.response.status_code = error_code
mocked_put.counter += 1
raise exc
else:
return put(*args, **kwargs)
def test_put_overwrite_with_downscope(
tmpdir, conn_cnx, db_parameters, is_public_test, from_path
):
"""Tests whether _force_put_overwrite and overwrite=true works as intended."""
if is_public_test:
pytest.xfail(
"Server need to update with merged change. Expected release version: 4.41.0"
)
with conn_cnx() as cnx:
tmp_dir = str(tmpdir.mkdir("data"))
test_data = os.path.join(tmp_dir, "data.txt")
with open(test_data, "w") as f:
f.write("test1,test2")
f.write("test3,test4")
cnx.cursor().execute("RM @~/test_put_overwrite")
try:
file_stream = None if from_path else open(test_data, "rb")
with cnx.cursor() as cur:
cur.execute("ALTER SESSION SET GCS_USE_DOWNSCOPED_CREDENTIAL = TRUE")
put(
cur,
test_data,
"~/test_put_overwrite",
from_path,
file_stream=file_stream,
)
data = cur.fetchall()
assert data[0][6] == "UPLOADED"
put(
cur,
test_data,
"~/test_put_overwrite",
from_path,
file_stream=file_stream,
)
data = cur.fetchall()
assert data[0][6] == "SKIPPED"
put(
cur,
test_data,
"~/test_put_overwrite",
from_path,
sql_options="OVERWRITE = TRUE",
file_stream=file_stream,
)
data = cur.fetchall()
assert data[0][6] == "UPLOADED"
ret = cnx.cursor().execute("LS @~/test_put_overwrite").fetchone()
assert "test_put_overwrite/data.txt" in ret[0]
assert "data.txt.gz" in ret[0]
finally:
if file_stream:
file_stream.close()
cnx.cursor().execute("RM @~/test_put_overwrite")
