Skip to content

Commit

Permalink
fix: Fix bug with no SqlRegistryConfig class (feast-dev#3586)
Browse files Browse the repository at this point in the history
* adding SqlRegistryConfig class

Signed-off-by: davidschuler-8451 <david.schuler@8451.com>

* refactor: move SqlRegistryConfig class to sql.py

Signed-off-by: davidschuler-8451 <david.schuler@8451.com>

* enabling SqlRegistry to accept RegistryConfig or SqlRegistryConfig

Signed-off-by: davidschuler-8451 <david.schuler@8451.com>

---------

Signed-off-by: davidschuler-8451 <david.schuler@8451.com>
  • Loading branch information
davidschuler-8451 authored and zerafachris committed Mar 5, 2024
1 parent f1bd947 commit bf9c12e
Showing 1 changed file with 20 additions and 16 deletions.
36 changes: 20 additions & 16 deletions sdk/python/feast/infra/registry/sql.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
from threading import Lock
from typing import Any, Callable, List, Optional, Set, Union

from pydantic import StrictStr
from sqlalchemy import ( # type: ignore
BigInteger,
Column,
Expand Down Expand Up @@ -164,28 +165,31 @@ def get_registry_tables(metadata : MetaData = MetaData()):
Column("infra_proto", LargeBinary, nullable=False),
)

feast_metadata = Table(
"feast_metadata",
metadata,
Column("project_id", String(50), primary_key=True),
Column("metadata_key", String(50), primary_key=True),
Column("metadata_value", String(50), nullable=False),
Column("last_updated_timestamp", BigInteger, nullable=False),
)

return entities, data_sources, feature_views, request_feature_views, \
stream_feature_views, on_demand_feature_views, feature_services, \
saved_datasets, validation_references, managed_infra, feast_metadata
feast_metadata = Table(
"feast_metadata",
metadata,
Column("project_id", String(50), primary_key=True),
Column("metadata_key", String(50), primary_key=True),
Column("metadata_value", String(50), nullable=False),
Column("last_updated_timestamp", BigInteger, nullable=False),
)


class SqlRegistryConfig(RegistryConfig):
registry_type: StrictStr = "sql"
""" str: Provider name or a class name that implements Registry."""

class FeastMetadataKeys(Enum):
LAST_UPDATED_TIMESTAMP = "last_updated_timestamp"
PROJECT_UUID = "project_uuid"
path: StrictStr = ""
""" str: Path to metadata store.
If registry_type is 'sql', then this is a database URL as expected by SQLAlchemy """


class SqlRegistry(BaseRegistry):
def __init__(
self, registry_config: Optional[RegistryConfig], project: str, repo_path: Optional[Path]
self,
registry_config: Optional[Union[RegistryConfig, SqlRegistryConfig]],
project: str,
repo_path: Optional[Path],
):
assert registry_config is not None, "SqlRegistry needs a valid registry_config"
self.engine: Engine = create_engine(registry_config.path, echo=False)
Expand Down

0 comments on commit bf9c12e

Please sign in to comment.