def publish(project: Project, event: threading.Event):
fake = Faker()
while not event.is_set():
rec = {
"name": fake.name(),
"phone": fake.phone_number()
}
project.send(rec)
time.sleep(1)
def test_send_bulk(reader, client: Client):
client._write_records = Mock()
records = [{"one": "test", "two": "test"}]
p = Project(name="proj", client=client, project_id=123)
p.send_bulk(
records, detection_mode="all", params={"test": "param"}, headers={"test": "two"}
)
client._write_records.assert_called_with(
project="proj",
reader=reader.return_value,
headers={"test": "two"},
params={"test": "param", "detection_mode": "all"},
)
def test_send(client: Client):
client._write_record_sync = Mock(
return_value={"data": {"success": [], "failure": []}}
)
records = [{"one": "test", "two": "test"}]
p = Project(name="proj", client=client, project_id=123)
p.send(
records, detection_mode="all", params={"test": "param"}, headers={"test": "two"}
)
client._write_record_sync.assert_called_with(
"proj",
records,
headers={"test": "two"},
params={"test": "param", "detection_mode": "all"},
)
def subscriber(project: Project):
# By default, ``iter_records()`` will block waiting for
# new labeled records to be ready for consumption. It will
# fetch in batches, so if there are already records in the
# stream, we will immediately yield the most recent 200 records.
for rec in project.iter_records():
yield rec
def _create_get_project(self, *, name: str, desc: str, display_name: str):
res = self._create_project(name=name,
desc=desc,
display_name=display_name)
_id = res["data"]["id"]
p = self._get_project(_id)["project"]
return Project(
name=p["name"],
client=self,
project_id=_id,
desc=p["description"],
display_name=p["display_name"],
)
def search_projects(self, count=200, query=None) -> List[Project]:
"""Search for projects that you are a member of.
Args:
count: the max number of projects to return
query: an optional query string to filter projects on
Returns:
A list of ``Project`` instances
"""
params = {"limit": count}
if query is not None:
params.update({"query": query})
projects = self._get("projects", params=params)["data"]["projects"]
return [
Project(name=p["name"],
client=self,
project_id=p["_id"],
desc=p["description"]) for p in projects
]
def get_project(
self,
*,
name: str = None,
create: bool = False,
desc: str = None,
display_name: str = None,
) -> Project:
"""
Create or get a project. By default, this method will try
and fetch an existing Gretel project. If the project does not
exist or you are not a member, a ``NotFound`` exception will be thrown.
Additionally, you can try and create the project by setting the ``create``
flag to ``True``. If this flag is set and the project already exists,
``Project`` instance will be returned. If the project does not already
exist, Gretel will attempt to create it for you. If the project name is
not available then a ``BadRequest`` will be returned.
Finally, if you just need a quick project to work with, Gretel can name
the project for you by omitting the ``name``::
client.get_project(create=True)
Args:
name: The unique name of the project to get or create.
create: If the project does not exist, try and create it. If no project name
is provided, create a unique name based on the authenticated user.
desc: If project gets created, set the description to this value. This will only
get used when a project gets newly created. If the project already exists,
nothing will happen with this value.
display_name: If project gets created, set the display name to this value. This will
be the primary name used when looking at the project in the Gretel Console.
Returns:
A ``Project`` instance.
Raises:
``Unauthorized`` or ``BadRequest``
"""
if not name and not create: # pragma: no cover
raise ValueError("Must provide a name or create flag!")
if name:
try:
p = self._get_project(name)["project"]
return Project(
name=name,
client=self,
project_id=p["_id"],
desc=p["description"],
display_name=p["display_name"],
)
except (NotFound, Unauthorized):
if create:
return self._create_get_project(name=name,
desc=desc,
display_name=display_name)
else:
raise
if not name and create:
return self._create_get_project(name=None,
desc=desc,
display_name=display_name)
def test_get_console_link(client: Client):
p = Project(name="test-project", client=client, project_id="1234")
assert p.get_console_url() == "https://console.gretel.cloud/test-project"
def test_send_dataframe(client: Client):
client._write_records = Mock()
df = pd.DataFrame([{f"foo_{i}": "bar"} for i in range(50)])
p = Project(name="proj", client=client, project_id=123)
p.send_dataframe(
df,
sample=25,
detection_mode="none",
headers={"X-Test-Gretel": "one"},
params={"test-param": "two"},
)
_, _, kwargs = client._write_records.mock_calls[0]
check_df = kwargs["reader"]
assert len(check_df.df) == 25
assert kwargs["headers"] == {"X-Test-Gretel": "one"}
assert kwargs["params"] == {"test-param": "two", "detection_mode": "none"}
with pytest.raises(ValueError):
p.send_dataframe(df, sample=0)
with pytest.raises(ValueError):
p.send_dataframe(df, sample=100)
with pytest.raises(ValueError):
p.send_dataframe(df, sample=-1)
with pytest.raises(ValueError):
p.send_dataframe([1, 2])
p.send_dataframe(df, sample=0.1)
_, _, kwargs = client._write_records.mock_calls[1]
check_df = kwargs["reader"]
assert len(check_df.df) == 5
p.send_dataframe(df)
_, _, kwargs = client._write_records.mock_calls[2]
check_df = kwargs["reader"]
assert len(check_df.df) == 50
def test_iter_records(client: Client):
client._iter_records = Mock()
p = Project(name="proj", client=client, project_id=123)
p.iter_records()
_, _, kwargs = p.client._iter_records.mock_calls[0]
assert kwargs["project"] == "proj"