1
1
import importlib
2
2
import inspect
3
+ import logging
4
+ import os
3
5
import pkgutil
4
6
import sys
5
7
from typing import Set
6
8
7
9
import setuptools
8
10
import typer
9
11
10
- from butterfree ._cli import cli_logger
12
+ from butterfree .clients import SparkClient
13
+ from butterfree .configs import environment
14
+ from butterfree .extract .readers import FileReader
11
15
from butterfree .migrations .database_migration import ALLOWED_DATABASE
12
16
from butterfree .pipelines import FeatureSetPipeline
13
17
14
- app = typer .Typer ()
18
+ app = typer .Typer (help = "Apply the automatic migrations in a database." )
19
+
20
+ logger = logging .getLogger ("migrate" )
15
21
16
22
17
23
def __find_modules (path : str ) -> Set [str ]:
@@ -33,18 +39,18 @@ def __find_modules(path: str) -> Set[str]:
33
39
34
40
35
41
def __fs_objects (path : str ) -> Set [FeatureSetPipeline ]:
36
- cli_logger .info (f"Looking for python modules under { path } ..." )
42
+ logger .info (f"Looking for python modules under { path } ..." )
37
43
modules = __find_modules (path )
38
44
if not modules :
39
45
return set ()
40
46
41
- cli_logger .info (f"Importing modules..." )
47
+ logger .info (f"Importing modules..." )
42
48
package = "." .join (path .strip ("/" ).split ("/" ))
43
49
imported = set (
44
50
importlib .import_module (f".{ name } " , package = package ) for name in modules
45
51
)
46
52
47
- cli_logger .info (f"Scanning modules..." )
53
+ logger .info (f"Scanning modules..." )
48
54
content = {
49
55
module : set (
50
56
filter (
@@ -80,14 +86,18 @@ def __fs_objects(path: str) -> Set[FeatureSetPipeline]:
80
86
81
87
instances .add (value )
82
88
83
- cli_logger .info ("Creating instances..." )
89
+ logger .info ("Creating instances..." )
84
90
return set (value () for value in instances )
85
91
86
92
87
93
PATH = typer .Argument (
88
94
..., help = "Full or relative path to where feature set pipelines are being defined." ,
89
95
)
90
96
97
+ GENERATE_LOGS = typer .Option (
98
+ False , help = "To generate the logs in local file 'logging.json'."
99
+ )
100
+
91
101
92
102
class Migrate :
93
103
"""Execute migration operations in a Database based on pipeline Writer.
@@ -96,23 +106,43 @@ class Migrate:
96
106
pipelines: list of Feature Set Pipelines to use to migration.
97
107
"""
98
108
99
- def __init__ (self , pipelines : Set [FeatureSetPipeline ]) -> None :
109
+ def __init__ (
110
+ self , pipelines : Set [FeatureSetPipeline ], spark_client : SparkClient = None
111
+ ) -> None :
100
112
self .pipelines = pipelines
113
+ self .spark_client = spark_client or SparkClient ()
101
114
102
- def _send_logs_to_s3 (self ) -> None :
115
+ def _send_logs_to_s3 (self , file_local : bool ) -> None :
103
116
"""Send all migration logs to S3."""
104
- pass
117
+ file_reader = FileReader (id = "name" , path = "logs/logging.json" , format = "json" )
118
+ df = file_reader .consume (self .spark_client )
119
+
120
+ path = environment .get_variable ("FEATURE_STORE_S3_BUCKET" )
105
121
106
- def run (self ) -> None :
122
+ self .spark_client .write_dataframe (
123
+ dataframe = df ,
124
+ format_ = "json" ,
125
+ mode = "append" ,
126
+ ** {"path" : f"s3a://{ path } /logging" },
127
+ )
128
+
129
+ if not file_local :
130
+ os .rmdir ("logs/logging.json" )
131
+
132
+ def run (self , generate_logs : bool ) -> None :
107
133
"""Construct and apply the migrations."""
108
134
for pipeline in self .pipelines :
109
135
for writer in pipeline .sink .writers :
110
136
migration = ALLOWED_DATABASE [writer .db_config .database ]
111
137
migration .apply_migration (pipeline .feature_set , writer )
112
138
139
+ self ._send_logs_to_s3 (generate_logs )
140
+
113
141
114
- @app .callback ()
115
- def migrate (path : str = PATH ) -> Set [FeatureSetPipeline ]:
142
+ @app .command ("apply" )
143
+ def migrate (
144
+ path : str = PATH , generate_logs : bool = GENERATE_LOGS ,
145
+ ) -> Set [FeatureSetPipeline ]:
116
146
"""Scan and run database migrations for feature set pipelines defined under PATH.
117
147
118
148
Butterfree will scan a given path for classes that inherit from its
@@ -124,5 +154,5 @@ def migrate(path: str = PATH) -> Set[FeatureSetPipeline]:
124
154
import and instantiate them.
125
155
"""
126
156
pipe_set = __fs_objects (path )
127
- Migrate (pipe_set ).run ()
157
+ Migrate (pipe_set ).run (generate_logs )
128
158
return pipe_set
0 commit comments