-
Notifications
You must be signed in to change notification settings - Fork 1.3k
Expand file tree
/
Copy pathbootstrap.py
More file actions
82 lines (66 loc) · 2.99 KB
/
bootstrap.py
File metadata and controls
82 lines (66 loc) · 2.99 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
import click
import psycopg
from feast.file_utils import replace_str_in_file
from feast.infra.utils.postgres.connection_utils import df_to_postgres_table
from feast.infra.utils.postgres.postgres_config import PostgreSQLConfig
def bootstrap():
# Bootstrap() will automatically be called from the init_repo() during `feast init`
import pathlib
from datetime import datetime, timedelta
from feast.driver_test_data import create_driver_hourly_stats_df
repo_path = pathlib.Path(__file__).parent.absolute() / "feature_repo"
config_file = repo_path / "feature_store.yaml"
end_date = datetime.now().replace(microsecond=0, second=0, minute=0)
start_date = end_date - timedelta(days=15)
driver_entities = [1001, 1002, 1003, 1004, 1005]
driver_df = create_driver_hourly_stats_df(driver_entities, start_date, end_date)
postgres_host = click.prompt("Postgres host", default="localhost")
postgres_port = click.prompt("Postgres port", default="5432")
postgres_database = click.prompt("Postgres DB name", default="postgres")
postgres_schema = click.prompt("Postgres schema", default="public")
postgres_user = click.prompt("Postgres user")
postgres_password = click.prompt("Postgres password", hide_input=True)
postgres_sslmode = click.prompt(
"Postgres sslmode (disable, allow, prefer, require, verify-ca, verify-full)",
default="require",
)
if click.confirm(
'Should I upload example data to Postgres (overwriting "feast_driver_hourly_stats" table)?',
default=True,
):
config = PostgreSQLConfig(
host=postgres_host,
port=int(postgres_port),
database=postgres_database,
db_schema=postgres_schema,
user=postgres_user,
password=postgres_password,
sslmode=postgres_sslmode,
)
db_connection = psycopg.connect(
conninfo=(
f"postgresql://{postgres_user}"
f":{postgres_password}"
f"@{postgres_host}"
f":{int(postgres_port)}"
f"/{postgres_database}"
),
sslmode=postgres_sslmode,
options=f"-c search_path={postgres_schema}",
)
with db_connection as conn, conn.cursor() as cur:
cur.execute('DROP TABLE IF EXISTS "feast_driver_hourly_stats"')
df_to_postgres_table(
config=config,
df=driver_df,
table_name="feast_driver_hourly_stats",
)
replace_str_in_file(config_file, "DB_HOST", postgres_host)
replace_str_in_file(config_file, "DB_PORT", postgres_port)
replace_str_in_file(config_file, "DB_NAME", postgres_database)
replace_str_in_file(config_file, "DB_SCHEMA", postgres_schema)
replace_str_in_file(config_file, "DB_USERNAME", postgres_user)
replace_str_in_file(config_file, "DB_PASSWORD", postgres_password)
replace_str_in_file(config_file, "DB_SSLMODE", postgres_sslmode)
if __name__ == "__main__":
bootstrap()