chore: Add user agents for GCP/AWS services (feast-dev#3087)

* chore: Add BigQuery user agent

Signed-off-by: Danny Chiao <[email protected]>

* add other agents

Signed-off-by: Danny Chiao <[email protected]>

* fix

Signed-off-by: Danny Chiao <[email protected]>

Signed-off-by: Danny Chiao <[email protected]>

Author: adchia

sdk/python/feast/infra/offline_stores/bigquery.py

@@ -44,14 +44,15 @@
 from feast.repo_config import FeastConfigBaseModel, RepoConfig
 
 from ...saved_dataset import SavedDatasetStorage
-from ...usage import log_exceptions_and_usage
+from ...usage import get_user_agent, log_exceptions_and_usage
 from .bigquery_source import (
     BigQueryLoggingDestination,
     BigQuerySource,
     SavedDatasetBigQueryStorage,
 )
 
 try:
+    from google.api_core import client_info as http_client_info
     from google.api_core.exceptions import NotFound
     from google.auth.exceptions import DefaultCredentialsError
     from google.cloud import bigquery
@@ -65,6 +66,10 @@
     raise FeastExtrasDependencyImportError("gcp", str(e))
 
 
+def get_http_client_info():
+    return http_client_info.ClientInfo(user_agent=get_user_agent())
+
+
 class BigQueryOfflineStoreConfig(FeastConfigBaseModel):
     """Offline store config for GCP BigQuery"""
 
@@ -688,7 +693,9 @@ def _get_bigquery_client(
     project: Optional[str] = None, location: Optional[str] = None
 ) -> bigquery.Client:
     try:
-        client = bigquery.Client(project=project, location=location)
+        client = bigquery.Client(
+            project=project, location=location, client_info=get_http_client_info()
+        )
     except DefaultCredentialsError as e:
         raise FeastProviderLoginError(
             str(e)

sdk/python/feast/infra/online_stores/datastore.py

@@ -35,11 +35,12 @@
 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, tracing_span
+from feast.usage import get_user_agent, log_exceptions_and_usage, tracing_span
 
 LOGGER = logging.getLogger(__name__)
 
 try:
+    from google.api_core import client_info as http_client_info
     from google.auth.exceptions import DefaultCredentialsError
     from google.cloud import datastore
     from google.cloud.datastore.client import Key
@@ -49,6 +50,10 @@
     raise FeastExtrasDependencyImportError("gcp", str(e))
 
 
+def get_http_client_info():
+    return http_client_info.ClientInfo(user_agent=get_user_agent())
+
+
 ProtoBatch = Sequence[
     Tuple[EntityKeyProto, Dict[str, ValueProto], datetime, Optional[datetime]]
 ]
@@ -331,8 +336,7 @@ def _initialize_client(
 ) -> datastore.Client:
     try:
         client = datastore.Client(
-            project=project_id,
-            namespace=namespace,
+            project=project_id, namespace=namespace, client_info=get_http_client_info()
         )
         return client
     except DefaultCredentialsError as e:

sdk/python/feast/infra/online_stores/dynamodb.py

@@ -30,10 +30,11 @@
 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, tracing_span
+from feast.usage import get_user_agent, log_exceptions_and_usage, tracing_span
 
 try:
     import boto3
+    from botocore.config import Config
     from botocore.exceptions import ClientError
 except ImportError as e:
     from feast.errors import FeastExtrasDependencyImportError
@@ -330,7 +331,12 @@ def _write_batch_non_duplicates(
 
 
 def _initialize_dynamodb_client(region: str, endpoint_url: Optional[str] = None):
-    return boto3.client("dynamodb", region_name=region, endpoint_url=endpoint_url)
+    return boto3.client(
+        "dynamodb",
+        region_name=region,
+        endpoint_url=endpoint_url,
+        config=Config(user_agent=get_user_agent()),
+    )
 
 
 def _initialize_dynamodb_resource(region: str, endpoint_url: Optional[str] = None):

sdk/python/feast/infra/utils/aws_utils.py

@@ -22,6 +22,7 @@
     RedshiftTableNameTooLong,
 )
 from feast.type_map import pa_to_athena_value_type, pa_to_redshift_value_type
+from feast.usage import get_user_agent
 
 try:
     import boto3
@@ -39,7 +40,10 @@ def get_redshift_data_client(aws_region: str):
     """
     Get the Redshift Data API Service client for the given AWS region.
     """
-    return boto3.client("redshift-data", config=Config(region_name=aws_region))
+    return boto3.client(
+        "redshift-data",
+        config=Config(region_name=aws_region, user_agent=get_user_agent()),
+    )
 
 
 def get_s3_resource(aws_region: str):
@@ -680,7 +684,9 @@ def get_athena_data_client(aws_region: str):
     """
     Get the athena Data API Service client for the given AWS region.
     """
-    return boto3.client("athena", config=Config(region_name=aws_region))
+    return boto3.client(
+        "athena", config=Config(region_name=aws_region, user_agent=get_user_agent())
+    )
 
 
 @retry(

sdk/python/feast/usage.py

@@ -54,6 +54,13 @@
     ).hexdigest(),
 }
 
+APPLICATION_NAME = "feast-dev/feast"
+USER_AGENT = "{}/{}".format(APPLICATION_NAME, get_version())
+
+
+def get_user_agent():
+    return USER_AGENT
+
 
 def set_current_project_uuid(project_uuid: str):
     _constant_attributes["project_id"] = project_uuid