Skip to content

Commit a1823b7

Browse files
woopwoop
authored
Flatten configuration structure for online store (feast-dev#1459)
* Refactor structure of default repo configuration Signed-off-by: Willem Pienaar <git@willem.co> * Fix types and default configuration Signed-off-by: Willem Pienaar <git@willem.co> * Remove scratch code Signed-off-by: Willem Pienaar <git@willem.co> * Fix telemetry tests Signed-off-by: Willem Pienaar <git@willem.co> * Add extra comment Signed-off-by: Willem Pienaar <git@willem.co> * Force configuration of FeatureStore Signed-off-by: Willem Pienaar <git@willem.co>
1 parent 614fa55 commit a1823b7

13 files changed

Lines changed: 258 additions & 219 deletions

‎sdk/python/feast/feature_store.py‎

Lines changed: 8 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -29,12 +29,7 @@
2929
)
3030
from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto
3131
from feast.registry import Registry
32-
from feast.repo_config import (
33-
LocalOnlineStoreConfig,
34-
OnlineStoreConfig,
35-
RepoConfig,
36-
load_repo_config,
37-
)
32+
from feast.repo_config import RepoConfig, load_repo_config
3833
from feast.telemetry import Telemetry
3934
from feast.version import get_version
4035

@@ -51,6 +46,12 @@ class FeatureStore:
5146
def __init__(
5247
self, repo_path: Optional[str] = None, config: Optional[RepoConfig] = None,
5348
):
49+
""" Initializes a new FeatureStore object. Used to manage a feature store.
50+
51+
Args:
52+
repo_path: Path to a `feature_store.yaml` used to configure the feature store
53+
config (RepoConfig): Configuration object used to configure the feature store
54+
"""
5455
self.repo_path = repo_path
5556
if repo_path is not None and config is not None:
5657
raise ValueError("You cannot specify both repo_path and config")
@@ -59,14 +60,7 @@ def __init__(
5960
elif repo_path is not None:
6061
self.config = load_repo_config(Path(repo_path))
6162
else:
62-
self.config = RepoConfig(
63-
registry="./registry.db",
64-
project="default",
65-
provider="local",
66-
online_store=OnlineStoreConfig(
67-
local=LocalOnlineStoreConfig(path="online_store.db")
68-
),
69-
)
63+
raise ValueError("Please specify one of repo_path or config")
7064

7165
registry_config = self.config.get_registry_config()
7266
self._registry = Registry(

‎sdk/python/feast/infra/gcp.py‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -28,9 +28,10 @@
2828
class GcpProvider(Provider):
2929
_gcp_project_id: Optional[str]
3030

31-
def __init__(self, config: Optional[DatastoreOnlineStoreConfig]):
32-
if config:
33-
self._gcp_project_id = config.project_id
31+
def __init__(self, config: RepoConfig):
32+
assert isinstance(config.online_store, DatastoreOnlineStoreConfig)
33+
if config and config.online_store and config.online_store.project_id:
34+
self._gcp_project_id = config.online_store.project_id
3435
else:
3536
self._gcp_project_id = None
3637

‎sdk/python/feast/infra/local.py‎

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import os
22
import sqlite3
33
from datetime import datetime
4+
from pathlib import Path
45
from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
56

67
import pandas as pd
@@ -20,16 +21,22 @@
2021
from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto
2122
from feast.protos.feast.types.Value_pb2 import Value as ValueProto
2223
from feast.registry import Registry
23-
from feast.repo_config import LocalOnlineStoreConfig, RepoConfig
24+
from feast.repo_config import RepoConfig, SqliteOnlineStoreConfig
2425

2526

2627
class LocalProvider(Provider):
2728
_db_path: str
2829

29-
def __init__(self, config: LocalOnlineStoreConfig):
30-
self._db_path = config.path
30+
def __init__(self, config: RepoConfig):
31+
32+
assert config is not None
33+
assert config.online_store is not None
34+
local_online_store_config = config.online_store
35+
assert isinstance(local_online_store_config, SqliteOnlineStoreConfig)
36+
self._db_path = local_online_store_config.path
3137

3238
def _get_conn(self):
39+
Path(self._db_path).parent.mkdir(exist_ok=True)
3340
return sqlite3.connect(
3441
self._db_path, detect_types=sqlite3.PARSE_DECLTYPES | sqlite3.PARSE_COLNAMES
3542
)

‎sdk/python/feast/infra/provider.py‎

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -127,15 +127,11 @@ def get_provider(config: RepoConfig) -> Provider:
127127
if config.provider == "gcp":
128128
from feast.infra.gcp import GcpProvider
129129

130-
return GcpProvider(
131-
config.online_store.datastore if config.online_store else None
132-
)
130+
return GcpProvider(config)
133131
elif config.provider == "local":
134132
from feast.infra.local import LocalProvider
135133

136-
assert config.online_store is not None
137-
assert config.online_store.local is not None
138-
return LocalProvider(config.online_store.local)
134+
return LocalProvider(config)
139135
else:
140136
raise ValueError(config)
141137

‎sdk/python/feast/repo_config.py‎

Lines changed: 65 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,9 @@
11
from pathlib import Path
2-
from typing import Optional, Union
32

43
import yaml
5-
from pydantic import BaseModel, StrictInt, StrictStr, ValidationError
4+
from pydantic import BaseModel, StrictInt, StrictStr, ValidationError, root_validator
5+
from pydantic.error_wrappers import ErrorWrapper
6+
from pydantic.typing import Dict, Literal, Optional, Union
67

78

89
class FeastBaseModel(BaseModel):
@@ -13,26 +14,27 @@ class Config:
1314
extra = "forbid"
1415

1516

16-
class LocalOnlineStoreConfig(FeastBaseModel):
17-
""" Online store config for local (SQLite-based) online store """
17+
class SqliteOnlineStoreConfig(FeastBaseModel):
18+
""" Online store config for local (SQLite-based) store """
1819

19-
path: StrictStr
20-
""" str: Path to sqlite db """
20+
type: Literal["sqlite"] = "sqlite"
21+
""" Online store type selector"""
22+
23+
path: StrictStr = "data/online.db"
24+
""" (optional) Path to sqlite db """
2125

2226

2327
class DatastoreOnlineStoreConfig(FeastBaseModel):
2428
""" Online store config for GCP Datastore """
2529

26-
project_id: StrictStr
27-
""" str: GCP Project Id """
30+
type: Literal["datastore"] = "datastore"
31+
""" Online store type selector"""
2832

33+
project_id: Optional[StrictStr] = None
34+
""" (optional) GCP Project Id """
2935

30-
class OnlineStoreConfig(FeastBaseModel):
31-
datastore: Optional[DatastoreOnlineStoreConfig] = None
32-
""" DatastoreOnlineStoreConfig: Optional Google Cloud Datastore config """
3336

34-
local: Optional[LocalOnlineStoreConfig] = None
35-
""" LocalOnlineStoreConfig: Optional local online store config """
37+
OnlineStoreConfig = Union[DatastoreOnlineStoreConfig, SqliteOnlineStoreConfig]
3638

3739

3840
class RegistryConfig(FeastBaseModel):
@@ -51,7 +53,7 @@ class RegistryConfig(FeastBaseModel):
5153
class RepoConfig(FeastBaseModel):
5254
""" Repo config. Typically loaded from `feature_store.yaml` """
5355

54-
registry: Union[StrictStr, RegistryConfig]
56+
registry: Union[StrictStr, RegistryConfig] = "data/registry.db"
5557
""" str: Path to metadata store. Can be a local path, or remote object storage path, e.g. gcs://foo/bar """
5658

5759
project: StrictStr
@@ -63,7 +65,7 @@ class RepoConfig(FeastBaseModel):
6365
provider: StrictStr
6466
""" str: local or gcp """
6567

66-
online_store: Optional[OnlineStoreConfig] = None
68+
online_store: OnlineStoreConfig = SqliteOnlineStoreConfig()
6769
""" OnlineStoreConfig: Online store configuration (optional depending on provider) """
6870

6971
def get_registry_config(self):
@@ -72,40 +74,54 @@ def get_registry_config(self):
7274
else:
7375
return self.registry
7476

75-
76-
# This is the JSON Schema for config validation. We use this to have nice detailed error messages
77-
# for config validation, something that bindr unfortunately doesn't provide out of the box.
78-
#
79-
# The schema should match the namedtuple structure above. It could technically even be inferred from
80-
# the types above automatically; but for now we choose a more tedious but less magic path of
81-
# providing the schema manually.
82-
83-
config_schema = {
84-
"type": "object",
85-
"properties": {
86-
"project": {"type": "string"},
87-
"registry": {"type": "string"},
88-
"provider": {"type": "string"},
89-
"online_store": {
90-
"type": "object",
91-
"properties": {
92-
"local": {
93-
"type": "object",
94-
"properties": {"path": {"type": "string"}},
95-
"additionalProperties": False,
96-
},
97-
"datastore": {
98-
"type": "object",
99-
"properties": {"project_id": {"type": "string"}},
100-
"additionalProperties": False,
101-
},
102-
},
103-
"additionalProperties": False,
104-
},
105-
},
106-
"required": ["project"],
107-
"additionalProperties": False,
108-
}
77+
@root_validator(pre=True)
78+
def _validate_online_store_config(cls, values):
79+
# This method will validate whether the online store configurations are set correctly. This explicit validation
80+
# is necessary because Pydantic Unions throw very verbose and cryptic exceptions. We also use this method to
81+
# impute the default online store type based on the selected provider. For the time being this method should be
82+
# considered tech debt until we can implement https://github.com/samuelcolvin/pydantic/issues/619 or a more
83+
# granular configuration system
84+
85+
# Skip if online store isn't set explicitly
86+
if "online_store" not in values:
87+
values["online_store"] = dict()
88+
89+
# Skip if we arent creating the configuration from a dict
90+
if not isinstance(values["online_store"], Dict):
91+
return values
92+
93+
# Make sure that the provider configuration is set. We need it to set the defaults
94+
assert "provider" in values
95+
96+
if "online_store" in values:
97+
# Set the default type
98+
if "type" not in values["online_store"]:
99+
if values["provider"] == "local":
100+
values["online_store"]["type"] = "sqlite"
101+
elif values["provider"] == "gcp":
102+
values["online_store"]["type"] = "datastore"
103+
104+
online_store_type = values["online_store"]["type"]
105+
106+
# Make sure the user hasn't provided the wrong type
107+
assert online_store_type in ["datastore", "sqlite"]
108+
109+
# Validate the dict to ensure one of the union types match
110+
try:
111+
if online_store_type == "sqlite":
112+
SqliteOnlineStoreConfig(**values["online_store"])
113+
elif values["online_store"]["type"] == "datastore":
114+
DatastoreOnlineStoreConfig(**values["online_store"])
115+
else:
116+
raise ValidationError(
117+
f"Invalid online store type {online_store_type}"
118+
)
119+
except ValidationError as e:
120+
raise ValidationError(
121+
[ErrorWrapper(e, loc="online_store")],
122+
model=SqliteOnlineStoreConfig,
123+
)
124+
return values
109125

110126

111127
class FeastConfigError(Exception):

‎sdk/python/feast/repo_operations.py‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -176,8 +176,7 @@ def init_repo(repo_path: Path, minimal: bool):
176176
registry: /path/to/registry.db
177177
provider: local
178178
online_store:
179-
local:
180-
path: /path/to/online_store.db
179+
path: /path/to/online_store.db
181180
"""
182181
)
183182
)
@@ -214,8 +213,7 @@ def init_repo(repo_path: Path, minimal: bool):
214213
registry: {"data/registry.db"}
215214
provider: local
216215
online_store:
217-
local:
218-
path: {"data/online_store.db"}
216+
path: {"data/online_store.db"}
219217
"""
220218
)
221219
)

‎sdk/python/telemetry_tests/test_telemetry.py‎

Lines changed: 50 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
14+
import tempfile
1415
import uuid
1516
from datetime import datetime
1617

@@ -21,7 +22,8 @@
2122
from time import sleep
2223
from importlib import reload
2324

24-
from feast import Client, Entity, ValueType, FeatureStore
25+
from feast import Client, Entity, ValueType, FeatureStore, RepoConfig
26+
from feast.repo_config import SqliteOnlineStoreConfig
2527

2628
TELEMETRY_BIGQUERY_TABLE = (
2729
"kf-feast.feast_telemetry.cloudfunctions_googleapis_com_cloud_functions"
@@ -91,19 +93,29 @@ def test_telemetry_on():
9193
os.environ["FEAST_IS_TELEMETRY_TEST"] = "True"
9294
os.environ["FEAST_TELEMETRY"] = "True"
9395

94-
test_feature_store = FeatureStore()
95-
entity = Entity(
96-
name="driver_car_id",
97-
description="Car driver id",
98-
value_type=ValueType.STRING,
99-
labels={"team": "matchmaking"},
100-
)
101-
102-
test_feature_store.apply([entity])
103-
104-
os.environ.clear()
105-
os.environ.update(old_environ)
106-
ensure_bigquery_telemetry_id_with_retry(test_telemetry_id)
96+
with tempfile.TemporaryDirectory() as temp_dir:
97+
test_feature_store = FeatureStore(
98+
config=RepoConfig(
99+
registry=os.path.join(temp_dir, "registry.db"),
100+
project="fake_project",
101+
provider="local",
102+
online_store=SqliteOnlineStoreConfig(
103+
path=os.path.join(temp_dir, "online.db")
104+
),
105+
)
106+
)
107+
entity = Entity(
108+
name="driver_car_id",
109+
description="Car driver id",
110+
value_type=ValueType.STRING,
111+
labels={"team": "matchmaking"},
112+
)
113+
114+
test_feature_store.apply([entity])
115+
116+
os.environ.clear()
117+
os.environ.update(old_environ)
118+
ensure_bigquery_telemetry_id_with_retry(test_telemetry_id)
107119

108120

109121
def test_telemetry_off():
@@ -113,20 +125,30 @@ def test_telemetry_off():
113125
os.environ["FEAST_TELEMETRY"] = "False"
114126
os.environ["FEAST_FORCE_TELEMETRY_UUID"] = test_telemetry_id
115127

116-
test_feature_store = FeatureStore()
117-
entity = Entity(
118-
name="driver_car_id",
119-
description="Car driver id",
120-
value_type=ValueType.STRING,
121-
labels={"team": "matchmaking"},
122-
)
123-
test_feature_store.apply([entity])
124-
125-
os.environ.clear()
126-
os.environ.update(old_environ)
127-
sleep(30)
128-
rows = read_bigquery_telemetry_id(test_telemetry_id)
129-
assert rows.total_rows == 0
128+
with tempfile.TemporaryDirectory() as temp_dir:
129+
test_feature_store = FeatureStore(
130+
config=RepoConfig(
131+
registry=os.path.join(temp_dir, "registry.db"),
132+
project="fake_project",
133+
provider="local",
134+
online_store=SqliteOnlineStoreConfig(
135+
path=os.path.join(temp_dir, "online.db")
136+
),
137+
)
138+
)
139+
entity = Entity(
140+
name="driver_car_id",
141+
description="Car driver id",
142+
value_type=ValueType.STRING,
143+
labels={"team": "matchmaking"},
144+
)
145+
test_feature_store.apply([entity])
146+
147+
os.environ.clear()
148+
os.environ.update(old_environ)
149+
sleep(30)
150+
rows = read_bigquery_telemetry_id(test_telemetry_id)
151+
assert rows.total_rows == 0
130152

131153

132154
@retry(wait=wait_exponential(multiplier=1, min=1, max=10), stop=stop_after_attempt(5))

0 commit comments

Comments
 (0)