Skip to content

Commit 5fc9e6e

Browse files
toping4445younggyu-oh
toping4445
authored and
younggyu-oh
committed
1.added test-python-universal-athena in Makefile 2.replaced database,bucket_name hardcoding to variable in AthenaDataSourceCreator
Signed-off-by: Youngkyu OH <toping4445@gmail.com>
1 parent 23905a1 commit 5fc9e6e

3 files changed

Lines changed: 57 additions & 19 deletions

File tree

  • Makefile
  • sdk/python
    • feast/infra/offline_stores/contrib/athena_offline_store/tests
    • tests/integration/feature_repos/universal/data_sources

‎Makefile‎

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -139,6 +139,27 @@ test-python-universal-trino:
139139
not test_universal_types" \
140140
sdk/python/tests
141141

142+
test-python-universal-athena:
143+
PYTHONPATH='.' \
144+
FULL_REPO_CONFIGS_MODULE=sdk.python.feast.infra.offline_stores.contrib.athena_repo_configuration \
145+
PYTEST_PLUGINS=feast.infra.offline_stores.contrib.athena_offline_store.tests \
146+
FEAST_USAGE=False IS_TEST=True \
147+
S3_DATABASE=sampledb \
148+
S3_BUCKET_NAME=sagemaker-yelo-test \
149+
python -m pytest -n 1 --integration \
150+
-k "not test_go_feature_server and \
151+
not test_logged_features_validation and \
152+
not test_lambda and \
153+
not test_feature_logging and \
154+
not test_offline_write and \
155+
not test_push_offline and \
156+
not test_historical_retrieval_with_validation and \
157+
not test_historical_features_persisting and \
158+
not test_historical_retrieval_fails_on_validation" \
159+
sdk/python/tests
160+
161+
162+
142163
test-python-universal-postgres:
143164
PYTHONPATH='.' \
144165
FULL_REPO_CONFIGS_MODULE=sdk.python.feast.infra.offline_stores.contrib.postgres_repo_configuration \
@@ -229,7 +250,7 @@ install-go-ci-dependencies:
229250
python -m pip install pybindgen==0.22.0 protobuf==3.20.1
230251

231252
install-protoc-dependencies:
232-
pip install grpcio-tools==1.47.0 mypy-protobuf==3.1.0
253+
pip install grpcio-tools==1.48.0 mypy-protobuf==3.1.0
233254

234255
compile-protos-go: install-go-proto-dependencies install-protoc-dependencies
235256
python setup.py build_go_protos

‎sdk/python/feast/infra/offline_stores/contrib/athena_offline_store/tests/data_source.py‎

Lines changed: 28 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import uuid
2+
import os
23
from typing import Any, Dict, List, Optional
34

45
import pandas as pd
@@ -31,12 +32,14 @@ def __init__(self, project_name: str, *args, **kwargs):
3132
super().__init__(project_name)
3233
self.client = aws_utils.get_athena_data_client("ap-northeast-2")
3334
self.s3 = aws_utils.get_s3_resource("ap-northeast-2")
34-
35+
data_source = os.environ.get("S3_DATA_SOURCE") if os.environ.get("S3_DATA_SOURCE") else "AwsDataCatalog"
36+
database = os.environ.get("S3_DATABASE") if os.environ.get("S3_DATABASE") else "default"
37+
bucket_name = os.environ.get("S3_BUCKET_NAME") if os.environ.get("S3_BUCKET_NAME") else "feast-integration-tests"
3538
self.offline_store_config = AthenaOfflineStoreConfig(
36-
data_source="AwsDataCatalog",
39+
data_source=f"{data_source}",
3740
region="ap-northeast-2",
38-
database="sampledb",
39-
s3_staging_location="s3://sagemaker-yelo-test/test_dir",
41+
database=f"{database}",
42+
s3_staging_location=f"s3://{bucket_name}/test_dir",
4043
)
4144

4245
def create_data_source(
@@ -49,22 +52,32 @@ def create_data_source(
4952
field_mapping: Dict[str, str] = None,
5053
) -> DataSource:
5154

52-
destination_name = self.get_prefixed_table_name(destination_name)
55+
table_name = destination_name
56+
s3_target = (
57+
self.offline_store_config.s3_staging_location
58+
+ "/"
59+
+ self.project_name
60+
+ "/"
61+
+ table_name
62+
+ "/"
63+
+ table_name
64+
+ ".parquet"
65+
)
5366

5467
aws_utils.upload_df_to_athena(
5568
self.client,
5669
self.offline_store_config.data_source,
5770
self.offline_store_config.database,
5871
self.s3,
59-
self.offline_store_config.s3_staging_location,
60-
destination_name,
72+
s3_target,
73+
table_name,
6174
df,
6275
)
6376

64-
self.tables.append(destination_name)
77+
self.tables.append(table_name)
6578

6679
return AthenaSource(
67-
table=destination_name,
80+
table=table_name,
6881
timestamp_field=timestamp_field,
6982
created_timestamp_column=created_timestamp_column,
7083
field_mapping=field_mapping or {"ts_1": "ts"},
@@ -78,7 +91,11 @@ def create_saved_dataset_destination(self) -> SavedDatasetAthenaStorage:
7891
)
7992
self.tables.append(table)
8093

81-
return SavedDatasetAthenaStorage(table_ref=table)
94+
return SavedDatasetAthenaStorage(
95+
table_ref=table,
96+
database=self.offline_store_config.database,
97+
data_source=self.offline_store_config.data_source,
98+
)
8299

83100
def create_logged_features_destination(self) -> LoggingDestination:
84101
table = self.get_prefixed_table_name(
@@ -92,7 +109,7 @@ def create_offline_store_config(self) -> FeastConfigBaseModel:
92109
return self.offline_store_config
93110

94111
def get_prefixed_table_name(self, suffix: str) -> str:
95-
return f"{self.project_name}.{suffix}"
112+
return f"{self.project_name}_{suffix}"
96113

97114
def teardown(self):
98115
for table in self.tables:

‎sdk/python/tests/integration/feature_repos/universal/data_sources/athena.py‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import uuid
2+
import os
23
from typing import Any, Dict, List, Optional
34

45
import pandas as pd
@@ -31,12 +32,14 @@ def __init__(self, project_name: str, *args, **kwargs):
3132
super().__init__(project_name)
3233
self.client = aws_utils.get_athena_data_client("ap-northeast-2")
3334
self.s3 = aws_utils.get_s3_resource("ap-northeast-2")
34-
35+
data_source = os.environ.get("S3_DATA_SOURCE") if os.environ.get("S3_DATA_SOURCE") else "AwsDataCatalog"
36+
database = os.environ.get("S3_DATABASE") if os.environ.get("S3_DATABASE") else "sampledb"
37+
bucket_name = os.environ.get("S3_BUCKET_NAME") if os.environ.get("S3_BUCKET_NAME") else "feast-integration-tests"
3538
self.offline_store_config = AthenaOfflineStoreConfig(
36-
data_source="AwsDataCatalog",
39+
data_source=f"{data_source}",
3740
region="ap-northeast-2",
38-
database="sampledb",
39-
s3_staging_location="s3://sagemaker-yelo-test/test_dir",
41+
database=f"{database}",
42+
s3_staging_location=f"s3://{bucket_name}/test_dir",
4043
)
4144

4245
def create_data_source(
@@ -49,9 +52,6 @@ def create_data_source(
4952
field_mapping: Dict[str, str] = None,
5053
) -> DataSource:
5154

52-
# destination_name = self.get_prefixed_table_name(destination_name)
53-
# test_name, table_name = destination_name.split('.')
54-
5555
table_name = destination_name
5656
s3_target = (
5757
self.offline_store_config.s3_staging_location

0 commit comments

Comments
 (0)