Change upsert methods to use server side upsert (#63)

Author: aksestok

VERSION

@@ -1 +1 @@
-2.5.0
+3.0.0

exabel_data_sdk/client/api/data_set_api.py

@@ -59,17 +59,28 @@ def create_data_set(self, data_set: DataSet) -> DataSet:
         response = self.client.create_data_set(CreateDataSetRequest(data_set=data_set.to_proto()))
         return DataSet.from_proto(response)
 
-    def update_data_set(self, data_set: DataSet, update_mask: FieldMask = None) -> DataSet:
+    def update_data_set(
+        self,
+        data_set: DataSet,
+        update_mask: FieldMask = None,
+        allow_missing: bool = False,
+    ) -> DataSet:
         """
         Update a data set.
 
         Args:
-            data_set:     The data set to update.
-            update_mask:  The fields to update. If not specified, the update behaves as a
-                          full update, overwriting all existing fields and properties.
+            data_set:       The data set to update.
+            update_mask:    The fields to update. If not specified, the update behaves as a
+                            full update, overwriting all existing fields and properties.
+            allow_missing:  If set to true, and the resource is not found, a new resource will be
+                            created. In this situation, the "update_mask" is ignored.
         """
         response = self.client.update_data_set(
-            UpdateDataSetRequest(data_set=data_set.to_proto(), update_mask=update_mask)
+            UpdateDataSetRequest(
+                data_set=data_set.to_proto(),
+                update_mask=update_mask,
+                allow_missing=allow_missing,
+            )
         )
         return DataSet.from_proto(response)
 

exabel_data_sdk/client/api/entity_api.py

@@ -132,49 +132,38 @@ def create_entity(self, entity: Entity, entity_type: str) -> Entity:
         )
         return Entity.from_proto(response)
 
-    def update_entity(self, entity: Entity, update_mask: FieldMask = None) -> Entity:
+    def update_entity(
+        self, entity: Entity, update_mask: FieldMask = None, allow_missing: bool = False
+    ) -> Entity:
         """
         Update an entity.
 
         Args:
             entity:         The entity to update.
             update_mask:    Fields to update. If not specified, the update behaves as a full update,
                             overwriting all existing fields and properties.
+            allow_missing:  If set to true, and the resource is not found, a new resource will be
+                            created. In this situation, the "update_mask" is ignored.
         """
         response = self.client.update_entity(
-            UpdateEntityRequest(entity=entity.to_proto(), update_mask=update_mask)
+            UpdateEntityRequest(
+                entity=entity.to_proto(),
+                update_mask=update_mask,
+                allow_missing=allow_missing,
+            )
         )
         return Entity.from_proto(response)
 
-    def upsert_entity(self, entity: Entity, assume_exists: bool = True) -> Entity:
+    def upsert_entity(self, entity: Entity) -> Entity:
         """
         Upsert an entity.
 
         If the entity already exists, update it by replacement. Otherwise, create it.
 
         Args:
-            entity:         The entity to upsert.
-            assume_exists:  If True, the entity is assumed to exist. Will try to to update
-                            the entity, and fall back to creating it if it does not exist,
-                            and vice versa.
-        """
-        if assume_exists:
-            try:
-                entity = self.update_entity(entity)
-            except RequestError as error:
-                if error.error_type == ErrorType.NOT_FOUND:
-                    entity = self.create_entity(entity, entity.get_entity_type())
-                else:
-                    raise error
-        else:
-            try:
-                entity = self.create_entity(entity, entity.get_entity_type())
-            except RequestError as error:
-                if error.error_type == ErrorType.ALREADY_EXISTS:
-                    entity = self.update_entity(entity)
-                else:
-                    raise error
-        return entity
+            entity: The entity to upsert.
+        """
+        return self.update_entity(entity, allow_missing=True)
 
     def delete_entity(self, name: str) -> None:
         """
@@ -230,8 +219,7 @@ def bulk_create_entities(
 
         def insert(entity: Entity) -> ResourceCreationStatus:
             if upsert:
-                # Upsert entities assuming they already exist.
-                self.upsert_entity(entity=entity, assume_exists=True)
+                self.upsert_entity(entity=entity)
                 return ResourceCreationStatus.UPSERTED
             # Optimistically insert the entity.
             # If the entity already exists, we'll get an ALREADY_EXISTS error from the backend,

exabel_data_sdk/client/api/relationship_api.py

@@ -91,7 +91,10 @@ def create_relationship_type(self, relationship_type: RelationshipType) -> Relat
         return RelationshipType.from_proto(response)
 
     def update_relationship_type(
-        self, relationship_type: RelationshipType, update_mask: FieldMask = None
+        self,
+        relationship_type: RelationshipType,
+        update_mask: FieldMask = None,
+        allow_missing: bool = False,
     ) -> RelationshipType:
         """
         Update a relationship type.
@@ -100,10 +103,14 @@ def update_relationship_type(
             relationship_type: The relationship type to update.
             update_mask:       The fields to update. If not specified, the update behaves as a
                                full update, overwriting all existing fields and properties.
+            allow_missing:     If set to true, and the resource is not found, a new resource will
+                               be created. In this situation, the "update_mask" is ignored.
         """
         response = self.client.update_relationship_type(
             UpdateRelationshipTypeRequest(
-                relationship_type=relationship_type.to_proto(), update_mask=update_mask
+                relationship_type=relationship_type.to_proto(),
+                update_mask=update_mask,
+                allow_missing=allow_missing,
             )
         )
         return RelationshipType.from_proto(response)
@@ -228,52 +235,37 @@ def create_relationship(self, relationship: Relationship) -> Relationship:
         return Relationship.from_proto(response)
 
     def update_relationship(
-        self, relationship: Relationship, update_mask: FieldMask = None
+        self, relationship: Relationship, update_mask: FieldMask = None, allow_missing: bool = False
     ) -> Relationship:
         """
         Update a relationship between two entities.
 
         Args:
-            relationship:  The relationship to update.
-            update_mask:   The fields to update. If not specified, the update behaves as a
-                           full update, overwriting all existing fields and properties.
+            relationship:   The relationship to update.
+            update_mask:    The fields to update. If not specified, the update behaves as a
+                            full update, overwriting all existing fields and properties.
+            allow_missing:  If set to true, and the resource is not found, a new resource will
+                            be created. In this situation, the "update_mask" is ignored.
         """
         response = self.client.update_relationship(
-            UpdateRelationshipRequest(relationship=relationship.to_proto(), update_mask=update_mask)
+            UpdateRelationshipRequest(
+                relationship=relationship.to_proto(),
+                update_mask=update_mask,
+                allow_missing=allow_missing,
+            )
         )
         return Relationship.from_proto(response)
 
-    def upsert_relationship(
-        self, relationship: Relationship, assume_exists: bool = True
-    ) -> Relationship:
+    def upsert_relationship(self, relationship: Relationship) -> Relationship:
         """
         Upsert a relationship between two entities.
 
         If the relationship already exists, update it by replacement. Otherwise, create it.
 
         Args:
             relationship:   The relationship to upsert.
-            assume_exists:  If True, the relationship is assumed to exist. Will try to to update
-                            the relationship, and fall back to creating it if it does not exist,
-                            and vice versa.
         """
-        if assume_exists:
-            try:
-                relationship = self.update_relationship(relationship)
-            except RequestError as error:
-                if error.error_type == ErrorType.NOT_FOUND:
-                    relationship = self.create_relationship(relationship)
-                else:
-                    raise error
-        else:
-            try:
-                relationship = self.create_relationship(relationship)
-            except RequestError as error:
-                if error.error_type == ErrorType.ALREADY_EXISTS:
-                    relationship = self.update_relationship(relationship)
-                else:
-                    raise error
-        return relationship
+        return self.update_relationship(relationship, allow_missing=True)
 
     def delete_relationship(self, relationship_type: str, from_entity: str, to_entity: str) -> None:
         """
@@ -319,8 +311,7 @@ def bulk_create_relationships(
 
         def insert(relationship: Relationship) -> ResourceCreationStatus:
             if upsert:
-                # Upsert relationships assuming they already exist.
-                self.upsert_relationship(relationship=relationship, assume_exists=True)
+                self.upsert_relationship(relationship=relationship)
                 return ResourceCreationStatus.UPSERTED
             # Optimistically insert the relationship.
             # If the relationship already exists, we'll get an ALREADY_EXISTS error from the

exabel_data_sdk/client/api/signal_api.py

@@ -77,17 +77,33 @@ def create_signal(self, signal: Signal, create_library_signal: bool = False) ->
         )
         return Signal.from_proto(response)
 
-    def update_signal(self, signal: Signal, update_mask: FieldMask = None) -> Signal:
+    def update_signal(
+        self,
+        signal: Signal,
+        update_mask: FieldMask = None,
+        allow_missing: bool = False,
+        create_library_signal: bool = False,
+    ) -> Signal:
         """
         Update one signal and return it.
 
         Args:
-            signal: The signal to update.
-            update_mask: The fields to update. If not specified, the update behaves as a
-                         full update, overwriting all existing fields and properties.
+            signal:         The signal to update.
+            update_mask:    The fields to update. If not specified, the update behaves as a
+                            full update, overwriting all existing fields and properties.
+            allow_missing:  If set to true, and the resource is not found, a new resource will be
+                            created. In this situation, the "update_mask" is ignored.
+            create_library_signal:
+                            If allow_missing is set to true and the signal does not exist, also add
+                            it to the library.
         """
         response = self.client.update_signal(
-            UpdateSignalRequest(signal=signal.to_proto(), update_mask=update_mask),
+            UpdateSignalRequest(
+                signal=signal.to_proto(),
+                update_mask=update_mask,
+                allow_missing=allow_missing,
+                create_library_signal=create_library_signal,
+            ),
         )
         return Signal.from_proto(response)
 

exabel_data_sdk/client/api/time_series_api.py

@@ -176,7 +176,7 @@ def upsert_time_series(
         series: pd.Series,
         create_tag: bool = False,
         default_known_time: DefaultKnownTime = None,
-    ) -> bool:
+    ) -> None:
         """
         Create or update a time series.
 
@@ -197,20 +197,10 @@ def upsert_time_series(
                         the Known Time for data points where a specific known time timestamp
                         has not been given. If not provided, the Exabel API defaults to the
                         current time (upload time) as the Known Time.
-
-        Returns:
-            True if the time series already existed, or False if it is created
         """
-        try:
-            # Optimistically assume that the time series exists, and append to it.
-            # If it doesn't exist, we catch the error below and create the time series instead.
-            self.append_time_series_data(name, series, default_known_time)
-            return True
-        except RequestError as error:
-            if error.error_type == ErrorType.NOT_FOUND:
-                self.create_time_series(name, series, create_tag, default_known_time)
-                return False
-            raise
+        self.append_time_series_data(
+            name, series, default_known_time, allow_missing=True, create_tag=create_tag
+        )
 
     def clear_time_series_data(self, name: str, start: pd.Timestamp, end: pd.Timestamp) -> None:
         """
@@ -227,7 +217,12 @@ def clear_time_series_data(self, name: str, start: pd.Timestamp, end: pd.Timesta
         )
 
     def append_time_series_data(
-        self, name: str, series: pd.Series, default_known_time: DefaultKnownTime = None
+        self,
+        name: str,
+        series: pd.Series,
+        default_known_time: DefaultKnownTime = None,
+        allow_missing: bool = False,
+        create_tag: bool = False,
     ) -> None:
         """
         Append data to the given time series.
@@ -236,25 +231,36 @@ def append_time_series_data(
         overwritten.
 
         Args:
-            name:   The resource name of the time series.
-            series: Series with data to append.
+            name:           The resource name of the time series.
+            series:         Series with data to append.
             default_known_time:
-                        Specify a default known time policy. This is used to determine
-                        the Known Time for data points where a specific known time timestamp
-                        has not been given. If not provided, the Exabel API defaults to the
-                        current time (upload time) as the Known Time.
+                            Specify a default known time policy. This is used to determine
+                            the Known Time for data points where a specific known time timestamp
+                            has not been given. If not provided, the Exabel API defaults to the
+                            current time (upload time) as the Known Time.
+            allow_missing:  If set to true, and the resource is not found, a new resource will be
+                            created. In this situation, the "update_mask" is ignored.
+            create_tag:     If allow_missing is set to true and the time series does not exist, also
+                            create a tag for every entity type the signal has time series for.
         """
         self.client.update_time_series(
             UpdateTimeSeriesRequest(
                 time_series=ProtoTimeSeries(
                     name=name, points=self._series_to_time_series_points(series)
                 ),
                 default_known_time=default_known_time,
+                allow_missing=allow_missing,
+                create_tag=create_tag,
             ),
         )
 
     def append_time_series_data_and_return(
-        self, name: str, series: pd.Series, default_known_time: DefaultKnownTime = None
+        self,
+        name: str,
+        series: pd.Series,
+        default_known_time: DefaultKnownTime = None,
+        allow_missing: Optional[bool] = False,
+        create_tag: Optional[bool] = False,
     ) -> pd.Series:
         """
         Append data to the given time series, and return the full series.
@@ -263,13 +269,17 @@ def append_time_series_data_and_return(
         overwritten.
 
         Args:
-            name:   The resource name of the time series.
-            series: Series with data to append.
+            name:           The resource name of the time series.
+            series:         Series with data to append.
             default_known_time:
-                        Specify a default known time policy. This is used to determine
-                        the Known Time for data points where a specific known time timestamp
-                        has not been given. If not provided, the Exabel API defaults to the
-                        current time (upload time) as the Known Time.
+                            Specify a default known time policy. This is used to determine
+                            the Known Time for data points where a specific known time timestamp
+                            has not been given. If not provided, the Exabel API defaults to the
+                            current time (upload time) as the Known Time.
+            allow_missing:  If set to true, and the resource is not found, a new resource will be
+                            created. In this situation, the "update_mask" is ignored.
+            create_tag:     If allow_missing is set to true and the time series does not exist, also
+                            create a tag for every entity type the signal has time series for.
 
         Returns:
             A series with all data for the given time series.
@@ -282,6 +292,8 @@ def append_time_series_data_and_return(
                 ),
                 view=TimeSeriesView(time_range=TimeRange()),
                 default_known_time=default_known_time,
+                allow_missing=allow_missing,
+                create_tag=create_tag,
             ),
         )
         return self._time_series_points_to_series(time_series.points, time_series.name)
@@ -337,10 +349,10 @@ def bulk_upsert_time_series(
         """
 
         def insert(ts: pd.Series) -> ResourceCreationStatus:
-            existed = self.upsert_time_series(
+            self.upsert_time_series(
                 str(ts.name), ts, create_tag=create_tag, default_known_time=default_known_time
             )
-            return ResourceCreationStatus.UPSERTED if existed else ResourceCreationStatus.CREATED
+            return ResourceCreationStatus.UPSERTED
 
         return bulk_insert(series, insert, threads=threads, retries=retries)