Adjust neptune-client after endpoint changes (#1863)

neptune-ai · Sep 2, 2024 · 975556d · 975556d
1 parent 7881182
commit 975556d
Show file tree

Hide file tree

Showing 6 changed files with 12 additions and 65 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -30,6 +30,7 @@
 - Series values DTO conversion reworked with protocol buffer support ([#1738](https://github.com/neptune-ai/neptune-client/pull/1738))
 - Series values fetching reworked with protocol buffer support ([#1744](https://github.com/neptune-ai/neptune-client/pull/1744))
 - Added support for enhanced field definitions querying ([#1751](https://github.com/neptune-ai/neptune-client/pull/1751))
+- Added support for `NQL` `MATCHES` operator ([#1863](https://github.com/neptune-ai/neptune-client/pull/1863))
 
 ### Fixes
 - Fixed `tqdm.notebook` import only in Notebook environment ([#1716](https://github.com/neptune-ai/neptune-client/pull/1716))

diff --git a/src/neptune/internal/backends/hosted_neptune_backend.py b/src/neptune/internal/backends/hosted_neptune_backend.py
@@ -55,7 +55,6 @@
     LeaderboardEntry,
     NextPage,
     QueryFieldDefinitionsResult,
-    QueryFieldsResult,
     StringField,
     StringSeriesField,
     StringSeriesValues,
@@ -1033,31 +1032,6 @@ def get_float_series_values(
         except HTTPNotFound:
             raise FetchAttributeNotFoundException(path_to_str(path))
 
-    @with_api_exceptions_handler
-    def query_fields_within_project(
-        self,
-        project_id: QualifiedName,
-        field_names_filter: Optional[List[str]] = None,
-        experiment_ids_filter: Optional[List[str]] = None,
-        next_page: Optional[NextPage] = None,
-    ) -> QueryFieldsResult:
-        pagination = {"nextPage": next_page.to_dto()} if next_page else {}
-        params = {
-            "projectIdentifier": project_id,
-            "query": {
-                **pagination,
-                "attributeNamesFilter": field_names_filter,
-                "experimentIdsFilter": experiment_ids_filter,
-            },
-            **DEFAULT_REQUEST_KWARGS,
-        }
-
-        try:
-            result = self.leaderboard_client.api.queryAttributesWithinProject(**params).response().result
-            return QueryFieldsResult.from_model(result)
-        except HTTPNotFound:
-            raise ProjectNotFound(project_id=project_id)
-
     @with_api_exceptions_handler
     def fetch_atom_attribute_values(
         self, container_id: str, container_type: ContainerType, path: List[str]
@@ -1092,16 +1066,18 @@ def _get_file_set_download_request(self, container_id: str, container_type: Cont
             raise FetchAttributeNotFoundException(path_to_str(path))
 
     @with_api_exceptions_handler
-    def _get_column_types(self, project_id: UniqueId, column: str, types: Optional[Iterable[str]] = None) -> List[Any]:
+    def _get_column_types(self, project_id: UniqueId, column: str) -> List[Any]:
         params = {
             "projectIdentifier": project_id,
-            "search": column,
-            "type": types,
-            "params": {},
+            "query": {
+                "attributeNameFilter": {"mustMatchRegexes": [column]},
+            },
             **DEFAULT_REQUEST_KWARGS,
         }
         try:
-            return self.leaderboard_client.api.searchLeaderboardAttributes(**params).response().result.entries
+            return (
+                self.leaderboard_client.api.queryAttributeDefinitionsWithinProject(**params).response().result.entries
+            )
         except HTTPNotFound as e:
             raise ProjectNotFound(project_id=project_id) from e
 
@@ -1124,6 +1100,8 @@ def search_leaderboard_entries(
 
         step_size = min(default_step_size, limit) if limit else default_step_size
 
+        columns = set(columns) | {sort_by} if columns else {sort_by}
+
         types_filter = list(map(lambda container_type: container_type.to_api(), types)) if types else None
         attributes_filter = {"attributeFilters": [{"path": column} for column in columns]} if columns else {}
 
@@ -1132,7 +1110,7 @@ def search_leaderboard_entries(
         elif sort_by == "sys/id":
             sort_by_column_type = FieldType.STRING.value
         else:
-            sort_by_column_type_candidates = self._get_column_types(project_id, sort_by, types_filter)
+            sort_by_column_type_candidates = self._get_column_types(project_id, sort_by)
             sort_by_column_type = _get_column_type_from_entries(sort_by_column_type_candidates, sort_by)
 
         try:

diff --git a/src/neptune/internal/backends/neptune_backend.py b/src/neptune/internal/backends/neptune_backend.py
@@ -42,7 +42,6 @@
     LeaderboardEntry,
     NextPage,
     QueryFieldDefinitionsResult,
-    QueryFieldsResult,
     StringField,
     StringSeriesField,
     StringSeriesValues,
@@ -347,12 +346,3 @@ def query_fields_definitions_within_project(
         experiment_ids_filter: Optional[List[str]] = None,
         next_page: Optional[NextPage] = None,
     ) -> QueryFieldDefinitionsResult: ...
-
-    @abc.abstractmethod
-    def query_fields_within_project(
-        self,
-        project_id: QualifiedName,
-        field_names_filter: Optional[List[str]] = None,
-        experiment_ids_filter: Optional[List[str]] = None,
-        next_page: Optional[NextPage] = None,
-    ) -> QueryFieldsResult: ...
diff --git a/src/neptune/internal/backends/neptune_backend_mock.py b/src/neptune/internal/backends/neptune_backend_mock.py
@@ -52,7 +52,6 @@
     LeaderboardEntry,
     NextPage,
     QueryFieldDefinitionsResult,
-    QueryFieldsResult,
     StringField,
     StringPointValue,
     StringSeriesField,
@@ -819,15 +818,3 @@ def query_fields_definitions_within_project(
             entries=[],
             next_page=NextPage(next_page_token=None, limit=0),
         )
-
-    def query_fields_within_project(
-        self,
-        project_id: QualifiedName,
-        field_names_filter: Optional[List[str]] = None,
-        experiment_ids_filter: Optional[List[str]] = None,
-        next_page: Optional[NextPage] = None,
-    ) -> QueryFieldsResult:
-        return QueryFieldsResult(
-            entries=[],
-            next_page=NextPage(next_page_token=None, limit=0),
-        )
diff --git a/src/neptune/internal/backends/nql.py b/src/neptune/internal/backends/nql.py
@@ -76,6 +76,7 @@ class NQLAttributeOperator(str, Enum):
     CONTAINS = "CONTAINS"
     GREATER_THAN = ">"
     LESS_THAN = "<"
+    MATCHES = "MATCHES"
 
 
 class NQLAttributeType(str, Enum):

diff --git a/src/neptune/internal/backends/offline_neptune_backend.py b/src/neptune/internal/backends/offline_neptune_backend.py
@@ -38,7 +38,6 @@
     LeaderboardEntry,
     NextPage,
     QueryFieldDefinitionsResult,
-    QueryFieldsResult,
     StringField,
     StringSeriesField,
     StringSeriesValues,
@@ -186,12 +185,3 @@ def query_fields_definitions_within_project(
         next_page: Optional[NextPage] = None,
     ) -> QueryFieldDefinitionsResult:
         raise NeptuneOfflineModeFetchException
-
-    def query_fields_within_project(
-        self,
-        project_id: QualifiedName,
-        field_names_filter: Optional[List[str]] = None,
-        experiment_ids_filter: Optional[List[str]] = None,
-        next_page: Optional[NextPage] = None,
-    ) -> QueryFieldsResult:
-        raise NeptuneOfflineModeFetchException