Skip to content

Commit

Permalink
chore: Add types hints and set created_timestamp in sql registry (#3033)
Browse files Browse the repository at this point in the history
* chore: Add types for some methods in the sql registry

Signed-off-by: Achal Shah <achals@gmail.com>

* created_timestamp

Signed-off-by: Achal Shah <achals@gmail.com>
  • Loading branch information
achals authored Aug 8, 2022
1 parent 1996596 commit 63d541d
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 5 deletions.
18 changes: 13 additions & 5 deletions sdk/python/feast/infra/registry_stores/sql.py
Original file line number Diff line number Diff line change
Expand Up @@ -688,15 +688,16 @@ def _apply_object(
self,
table: Table,
project: str,
id_field_name,
obj,
proto_field_name,
name=None,
id_field_name: str,
obj: Any,
proto_field_name: str,
name: Optional[str] = None,
):
self._maybe_init_project_metadata(project)

name = name or obj.name if hasattr(obj, "name") else None
assert name, f"name needs to be provided for {obj}"

with self.engine.connect() as conn:
update_datetime = datetime.utcnow()
update_time = int(update_datetime.timestamp())
Expand All @@ -721,9 +722,16 @@ def _apply_object(
)
conn.execute(update_stmt)
else:
obj_proto = obj.to_proto()

if hasattr(obj_proto, "meta") and hasattr(
obj_proto.meta, "created_timestamp"
):
obj_proto.meta.created_timestamp.FromDatetime(update_datetime)

values = {
id_field_name: name,
proto_field_name: obj.to_proto().SerializeToString(),
proto_field_name: obj_proto.SerializeToString(),
"last_updated_timestamp": update_time,
"project_id": project,
}
Expand Down
6 changes: 6 additions & 0 deletions sdk/python/tests/unit/test_sql_registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -168,6 +168,9 @@ def test_apply_entity_success(sql_registry):
and entity.tags["team"] == "matchmaking"
)

# After the first apply, the created_timestamp should be the same as the last_update_timestamp.
assert entity.created_timestamp == entity.last_updated_timestamp

sql_registry.delete_entity("driver_car_id", project)
assert_project_uuid(project, project_uuid, sql_registry)
entities = sql_registry.list_entities(project)
Expand Down Expand Up @@ -256,6 +259,9 @@ def test_apply_feature_view_success(sql_registry):
and feature_view.entities[0] == "fs1_my_entity_1"
)

# After the first apply, the created_timestamp should be the same as the last_update_timestamp.
assert feature_view.created_timestamp == feature_view.last_updated_timestamp

sql_registry.delete_feature_view("my_feature_view_1", project)
feature_views = sql_registry.list_feature_views(project)
assert len(feature_views) == 0
Expand Down

0 comments on commit 63d541d

Please sign in to comment.