feat: ingestion support for The Atlas of Economic Complexity

.env.example

@@ -18,6 +18,7 @@ S3_BUCKET=lakehouse
 S3_INGEST_PREFIX=raw
 S3_STAGE_PREFIX=stage
 S3_GRAPHS_MART_PREFIX=marts/graphs
+S3_ANALYTICS_MART_PREFIX=marts/analytics
 S3_EXPORTS_PREFIX=exports
 S3_BACKUPS_PREFIX=backups
 
@@ -30,11 +31,13 @@ S3_BACKUPS_PREFIX=backups
 ENGINE_DB=engine.duckdb
 STAGE_DB=stage.sqlite
 GRAPHS_MART_DB=marts/graphs.sqlite
+ANALYTICS_MART_DB=marts/analytics.sqlite
 
 # KùzuDB configurations
 # =====================
 
-MUSIC_TASTE_GRAPH_DB=graphs/music_taste
+MUSIC_TASTE_GRAPH_DB=graphs/music_taste.kz
+ECON_COMP_GRAPH_DB=graphs/econ_comp.kz
 
 # Ollama configurations
 # =====================

dlctl/cli.py

@@ -13,6 +13,7 @@
 from export.cli import export
 from graph.cli import graph
 from ingest.cli import ingest
+from shared.cache import cache_usage, expunge_cache
 from shared.settings import LOCAL_DIR, MART_DB_VARS, env
 from shared.storage import Storage, StoragePrefix
 
@@ -185,7 +186,7 @@ def backup_ls(include_all: bool):
     help="Model name to transform (can be used multiple times)",
 )
 @click.option("--debug", is_flag=True, help="Run dbt with the debug flag")
-def transform(models: tuple[str], debug: bool):
+def transform(models: Optional[tuple[str, ...]], debug: bool):
     dbt_handler = DBTHandler(debug=debug)
     dbt_handler.run(models)
 
@@ -195,9 +196,18 @@ def transform(models: tuple[str], debug: bool):
 
 
 @dlctl.command(name="test", help="Run data tests")
-def test():
-    dbt_handler = DBTHandler()
-    dbt_handler.test()
+@click.option(
+    "--model",
+    "-m",
+    "models",
+    multiple=True,
+    type=click.STRING,
+    help="Model name to transform (can be used multiple times)",
+)
+@click.option("--debug", is_flag=True, help="Run dbt with the debug flag")
+def test(models: Optional[tuple[str, ...]], debug: bool):
+    dbt_handler = DBTHandler(debug=debug)
+    dbt_handler.test(models)
 
 
 # Documentation
@@ -243,5 +253,39 @@ def generate_init_sql(path: str):
     T.generate_init_sql(path)
 
 
+# Cache
+# =====
+
+
+@dlctl.group(help="Manage cache (requests, etc.)")
+def cache():
+    pass
+
+
+@cache.command(name="clean", help="Expunge cache")
+@click.option(
+    "-ns",
+    "--namespace",
+    type=click.Choice(["requests", "huggingface"]),
+    help="Limit cache cleaning to a namespace",
+)
+@click.option(
+    "-n",
+    "--name",
+    type=click.STRING,
+    help="Limit cache cleaning to a specific name (namespace required as well)",
+)
+def cache_clean(namespace: Optional[str], name: Optional[str]):
+    if namespace is None and name is not None:
+        raise click.UsageError("name requires that namespace is set")
+
+    expunge_cache(namespace, name)
+
+
+@cache.command(name="df", help="Calculate cache usage statistics")
+def cache_df():
+    cache_usage()
+
+
 if __name__ == "__main__":
     dlctl()

dlctl/dbt_handler.py

@@ -46,7 +46,7 @@ def mkdirs(self):
     def deps(self):
         self.dbt.invoke(["deps"] + self.PROJECT_ARGS)
 
-    def run(self, models: Optional[tuple[str]] = None):
+    def run(self, models: Optional[tuple[str, ...]] = None):
         args = ["run"]
         args += self.PROJECT_ARGS
 
@@ -71,8 +71,20 @@ def run(self, models: Optional[tuple[str]] = None):
             else:
                 log.warning("{}: {}", r.node.name, r.status)
 
-    def test(self):
-        self.dbt.invoke(["test"] + self.PROJECT_ARGS)
+    def test(self, models: Optional[tuple[str, ...]] = None):
+        args = ["test"]
+        args += self.PROJECT_ARGS
+
+        if self.debug:
+            args += ["--debug"]
+
+        if models is not None and len(models) > 0:
+            args += [
+                "--select",
+                ",".join(f"{model}" for model in models),
+            ]
+
+        self.dbt.invoke(args)
 
     def docs_generate(self):
         self.dbt.invoke(["docs", "generate"] + self.PROJECT_ARGS)

graph/cli.py

@@ -43,7 +43,14 @@ def load(schema: str, overwrite: bool):
 
     try:
         ops = KuzuOps(schema, overwrite=overwrite)
-        ops.load_music_graph(s3_path)
+
+        match schema:
+            case "music_taste":
+                ops.load_music_taste(s3_path)
+            case "econ_comp":
+                ops.load_econ_comp(s3_path)
+            case _:
+                raise click.UsageError(f"{schema}: graph unsupported")
     except Exception as e:
         log.error(e)