From 864bf97cb3126d4f6a76113432b320fa1ae84dda Mon Sep 17 00:00:00 2001
From: Ada Draginda <ada@makenotion.com>
Date: Fri, 3 Mar 2023 14:43:33 -0800
Subject: [PATCH] [bug] stop quering for entities once zero are returned (#10)

---
 pyproject.toml               | 2 +-
 src/datahub_tools/classes.py | 7 +++++++
 src/datahub_tools/client.py  | 2 ++
 3 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index 49f62a9..78417f9 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -5,7 +5,7 @@ build-backend = 'setuptools.build_meta'
 [project]
 name = 'datahub_tools'
 description = 'Python tools for working with DataHub'
-version = '1.0.1'
+version = '1.0.2'
 readme = 'README.md'
 requires-python = '>=3.7'
 dependencies = [
diff --git a/src/datahub_tools/classes.py b/src/datahub_tools/classes.py
index a03a701..5130093 100644
--- a/src/datahub_tools/classes.py
+++ b/src/datahub_tools/classes.py
@@ -5,6 +5,7 @@
 from dataclasses import dataclass
 from typing import Any, Dict, List, Optional
 
+import datahub.emitter.mce_builder as builder
 import jmespath
 
 
@@ -68,6 +69,12 @@ def has_tags(self) -> bool:
     def has_owners(self) -> bool:
         return bool(self.owners)
 
+    def get_platform(self):
+        return builder.dataset_urn_to_key(self.urn).platform
+
+    def get_name(self):
+        return builder.dataset_urn_to_key(self.urn).name
+
     def get_priority(self) -> Optional[str]:
         """
         return the priority of the dataset, as reflected by its tags (priority tags look like `Priority: P0`)
diff --git a/src/datahub_tools/client.py b/src/datahub_tools/client.py
index e3bc5f3..03cb576 100644
--- a/src/datahub_tools/client.py
+++ b/src/datahub_tools/client.py
@@ -271,6 +271,8 @@ def get_datahub_entities(
             else:
                 raise e
         else:
+            if not dh_entities:
+                break
             _start += _chunk_size
             # important: you can get more than one URN per qualified name because there
             # may be more than one platform (e.g. dbt, snowflake, etc.).