Add local Feast ingest image builds

metadata-ingestion/src/datahub/ingestion/source/feast.py

@@ -51,6 +51,7 @@ class FeastConfig(ConfigModel):
     # https://pymongo.readthedocs.io/en/stable/examples/authentication.html
     core_url: str = "localhost:6565"
     env: str = DEFAULT_ENV
+    use_local_build: bool = False
 
 
 @dataclass
@@ -101,8 +102,18 @@ def get_workunits(self) -> Iterable[MetadataWorkUnit]:
         with tempfile.NamedTemporaryFile(suffix=".json") as tf:
 
             docker_client = docker.from_env()
+
+            # image to use for initial feast extraction
+            feast_image = "feast-ingest"
+
+            # build the image locally if specified
+            if self.config.use_local_build:
+                image, _ = docker_client.images.build(path="../feast_image/")
+
+                feast_image = image.id
+
             docker_client.containers.run(
-                "feast-ingest",
+                feast_image,
                 f'python3 ingest.py --core_url="{self.config.core_url}" --output_path=/out.json',
                 # allow the image to access the core URL if on host
                 network_mode="host",
@@ -132,6 +143,10 @@ def get_workunits(self) -> Iterable[MetadataWorkUnit]:
                     MLFeaturePropertiesClass(
                         name=feature["name"],
                         dataType=self.get_field_type(feature["type"], feature["name"]),
+                        # TODO: update this
+                        sourceDataset=builder.make_dataset_urn(
+                            "test", "test", self.config.env
+                        ),
                     )
                 )
 

metadata-ingestion/src/datahub/ingestion/source/feast_image/.gitignore

@@ -0,0 +1,2 @@
+# ignore Feast image outputs
+mlruns

metadata-ingestion/tests/integration/feast/docker-compose.yml

@@ -20,15 +20,15 @@ services:
       - --spring.config.location=classpath:/application.yml,file:/etc/feast/application.yml
 
   setup:
-    image: feast-ingest
+    # build from the same Feast image used for ingestion
+    build: ../../../src/datahub/ingestion/source/feast_image/
     volumes:
       - ./make_tests.py:/app/make_tests.py
       - ./wait-for-it.sh:/app/wait-for-it.sh
     depends_on:
       - core
-    command:
-      ./wait-for-it.sh testfeast:6565 -- python3 make_tests.py
-      # - ls
+    # wait for Feast to start, then insert test data
+    command: ./wait-for-it.sh testfeast:6565 -- python3 make_tests.py
 
   db:
     image: postgres:12-alpine