def test_iter_record_with_limit(records):
client = get_cloud_client("api", "abc123xyz")
client._get = Mock(side_effect=records)
check = []
for rec in client._iter_records(project="foo", record_limit=3):
check.append(rec)
assert len(check) == 3
def test_iter_records_reverse(records_rev):
records = records_rev
client = get_cloud_client("api", "abcd123xyz")
client._get = Mock(side_effect=records)
check = []
for rec in client._iter_records(project="foo", direction="backward"):
check.append(rec["record"])
assert check == ["foo_1", "foo_2", "foo_3", "foo_4", "foo_5", "foo_6", "foo_7"]
def test_iter_records(records):
client = get_cloud_client("api", "abc123xyz")
client._get = Mock(side_effect=records)
check = []
for rec in client._iter_records(project="foo"):
check.append(rec["record"])
if len(check) == 6:
break
assert check == ["foo_3", "foo_2", "foo_1", "foo_4", "foo_6", "foo_5"]
def start(api_key: str):
client = get_cloud_client("api", api_key)
with temporary_project(client) as project:
publish_event = threading.Event()
publish_thread = threading.Thread(target=publish, args=(project, publish_event))
publish_thread.start()
try:
for rec in subscriber(project):
print(rec)
except KeyboardInterrupt:
print("Shutting down...")
finally:
publish_event.set()
publish_thread.join()
def main():
"""`gretel` script entrypoint"""
parser = parse_command()
command = parser.parse_args()
if not GRETEL_API_KEY and not command.api_key:
raise Exception("Gretel API key not set.")
gretel_client = get_cloud_client(GRETEL_ENV, command.api_key)
try:
command.func(command, gretel_client)
except AttributeError:
parser.print_help()
parser.exit()
def test_iter_records_reverse_single_record(records_rev):
records = records_rev
client = get_cloud_client("api", "abcd123xyz")
_rec = {"data": {"records": [{
"id": "7",
"data": "foo_7",
"metadata": {},
"ingest_time": "2020-05-10T12:41:55.585538",
}]}}
client._get = Mock(side_effect=[_rec, _rec])
check = []
for rec in client._iter_records(project="foo", direction="backward"):
check.append(rec["record"])
assert check == ["foo_7"]
def test_iter_entity_stream_custom_params(records):
client = get_cloud_client("api", "abc123xyz")
client._get = Mock(side_effect=records)
resp = list(
client._iter_records(
project="foo",
entity_stream="test_entity",
wait_for=1,
headers={"X-Gretel-Test": "true"},
)
)
assert client._get.called_with(
"streams/records/foo",
params={"entity_stream": "test_entity", "flatten": "yes", "with_meta": "yes"},
headers={"X-Gretel-Test": "true"},
)
def test_get_cloud_client_prompt(getenv, getpass, Client):
# when no env is set and prompt is true, ask for gretel key
getenv.return_value = None
get_cloud_client("api", "prompt")
assert getpass.call_count == 1
# when api key is set, and prompt is true, use api key
getenv.return_value = "abcd123"
get_cloud_client("api", "prompt")
Client.assert_called_with(host="api.gretel.cloud", api_key="abcd123")
assert getpass.call_count == 1
# when api key is set and prompt always is true, ask for api key
get_cloud_client("api", "prompt_always")
assert getpass.call_count == 2
# use api key env variable
get_cloud_client("api", "abc123")
Client.assert_called_with(host="api.gretel.cloud", api_key="abc123")
def client():
return get_cloud_client("api", "abcd123xyz")
def test_iter_records_does_terminate(records):
client = get_cloud_client("api", "abc123xyz")
client._get = Mock(side_effect=records)
for _ in client._iter_records(project="foo", wait_for=1):
continue
# -*- coding: utf-8 -*-
"""
Created on Tue Feb 2 01:05:54 2021
@author: nijum
"""
import pandas as pd
from gretel_client import get_cloud_client
pd.set_option('max_colwidth', None)
client = get_cloud_client(
prefix="api",
api_key=
"grtucbfbc5356cc1d49757a29f3a4c7eff225b2a4bf907b052fab9acc3b9e1bfd5b4")
client.install_packages()
import pandas as pd
dataset_path = 'https://gretel-public-website.s3-us-west-2.amazonaws.com/datasets/healthcare-analytics-vidhya/train_data.csv'
nrows = 10000 # We will use this later when generating data
training_df = pd.read_csv(dataset_path, nrows=nrows)
training_df
from pathlib import Path
checkpoint_dir = str(Path.cwd() / "checkpoints")
config_template = {
"checkpoint_dir": checkpoint_dir,
"vocab_size": 20000,
"overwrite": True
def client():
return get_cloud_client('api', 'abcd123xyz')
