chore: Update new contrib plugins to include usage (#2638)

adchia · web-flow · commit f22430a302f9 · 2022-05-04T14:35:22.000-07:00
Signed-off-by: Danny Chiao &lt;danny@tecton.ai&gt;
diff --git a/sdk/python/feast/infra/offline_stores/contrib/postgres_offline_store/postgres.py b/sdk/python/feast/infra/offline_stores/contrib/postgres_offline_store/postgres.py
@@ -40,6 +40,7 @@
 from feast.repo_config import RepoConfig
 from feast.saved_dataset import SavedDatasetStorage
 from feast.type_map import pg_type_code_to_arrow
+from feast.usage import log_exceptions_and_usage
 
 from .postgres_source import PostgreSQLSource
 
@@ -50,6 +51,7 @@ class PostgreSQLOfflineStoreConfig(PostgreSQLConfig):
 
 class PostgreSQLOfflineStore(OfflineStore):
     @staticmethod
+    @log_exceptions_and_usage(offline_store="postgres")
     def pull_latest_from_table_or_query(
         config: RepoConfig,
         data_source: DataSource,
@@ -100,6 +102,7 @@ def pull_latest_from_table_or_query(
         )
 
     @staticmethod
+    @log_exceptions_and_usage(offline_store="postgres")
     def get_historical_features(
         config: RepoConfig,
         feature_views: List[FeatureView],
@@ -186,6 +189,7 @@ def query_generator() -> Iterator[str]:
         )
 
     @staticmethod
+    @log_exceptions_and_usage(offline_store="postgres")
     def pull_all_from_table_or_query(
         config: RepoConfig,
         data_source: DataSource,
diff --git a/sdk/python/feast/infra/online_stores/contrib/hbase_online_store/hbase.py b/sdk/python/feast/infra/online_stores/contrib/hbase_online_store/hbase.py
@@ -14,6 +14,7 @@
 from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto
 from feast.protos.feast.types.Value_pb2 import Value as ValueProto
 from feast.repo_config import FeastConfigBaseModel, RepoConfig
+from feast.usage import log_exceptions_and_usage
 
 
 class HbaseOnlineStoreConfig(FeastConfigBaseModel):
@@ -78,6 +79,7 @@ def _get_conn(self, config: RepoConfig):
             self._conn = Connection(host=store_config.host, port=int(store_config.port))
         return self._conn
 
+    @log_exceptions_and_usage(online_store="hbase")
     def online_write_batch(
         self,
         config: RepoConfig,
@@ -128,6 +130,7 @@ def online_write_batch(
             b.put(row_key, values_dict)
         b.send()
 
+    @log_exceptions_and_usage(online_store="hbase")
     def online_read(
         self,
         config: RepoConfig,
@@ -142,6 +145,7 @@ def online_read(
             config: The RepoConfig for the current FeatureStore.
             table: Feast FeatureView.
             entity_keys: a list of entity keys that should be read from the FeatureStore.
+            requested_features: a list of requested feature names.
         """
         hbase = HbaseUtils(self._get_conn(config))
         project = config.project
@@ -172,6 +176,7 @@ def online_read(
                 result.append((res_ts, res))
         return result
 
+    @log_exceptions_and_usage(online_store="hbase")
     def update(
         self,
         config: RepoConfig,
diff --git a/sdk/python/feast/infra/online_stores/contrib/postgres.py b/sdk/python/feast/infra/online_stores/contrib/postgres.py
@@ -18,6 +18,7 @@
 from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto
 from feast.protos.feast.types.Value_pb2 import Value as ValueProto
 from feast.repo_config import RepoConfig
+from feast.usage import log_exceptions_and_usage
 
 
 class PostgreSQLOnlineStoreConfig(PostgreSQLConfig):
@@ -33,6 +34,7 @@ def _get_conn(self, config: RepoConfig):
             self._conn = _get_conn(config.online_store)
         return self._conn
 
+    @log_exceptions_and_usage(online_store="postgres")
     def online_write_batch(
         self,
         config: RepoConfig,
@@ -86,6 +88,7 @@ def online_write_batch(
                 if progress:
                     progress(len(cur_batch))
 
+    @log_exceptions_and_usage(online_store="postgres")
     def online_read(
         self,
         config: RepoConfig,
@@ -136,6 +139,7 @@ def online_read(
 
         return result
 
+    @log_exceptions_and_usage(online_store="postgres")
     def update(
         self,
         config: RepoConfig,
