little morning programming

scripts/generate_duckdb.py

@@ -4,6 +4,7 @@ from pathlib import Path
 import click
 from tqdm import tqdm
 from posix import mkdir
+from datetime import datetime, date
 
 sql_tmpl = jinja2.Template('''
 LOAD spatial;
@@ -14,6 +15,7 @@ CREATE TABLE IF NOT EXISTS city_road(
     points DOUBLE[][],
     tags MAP(VARCHAR, VARCHAR),
     highway TEXT,
+    extract_date Date,
     length DOUBLE
 );
 
@@ -59,7 +61,6 @@ CREATE TEMPORARY VIEW road_points_with_distance AS (
 
 CREATE TEMPORARY VIEW data AS (
     SELECT
-        '{{ city }}' AS city_name,
         relation_id AS relation_id,
         LIST([road_points_with_distance.lat, road_points_with_distance.lon]) AS points,
         ANY_VALUE(tags) AS tags,
@@ -71,8 +72,15 @@ CREATE TEMPORARY VIEW data AS (
 
 
 
-INSERT INTO city_road(city_name, relation_id, points, tags, highway, length) (
-    SELECT city_name, relation_id, points, tags, highway, length FROM data
+INSERT INTO city_road(city_name, extract_date, relation_id, points, tags, highway, length) (
+    SELECT '{{ city }}' AS city_name,
+           '{{extract_date}}' AS extract_date,
+           relation_id,
+           points,
+           tags,
+           highway,
+           length
+    FROM data
 );
 
 ''')
@@ -86,19 +94,23 @@ def sql_safe(w: str) -> str:
     return str(w).replace("'", r"\\'")
 
 
-def append_data_to_duckdb(conn: duckdb.DuckDBPyConnection, osm_pbf_path: Path) -> None:
+def append_data_to_duckdb(conn: duckdb.DuckDBPyConnection, osm_pbf_path: Path, extract_date: date) -> None:
     city = extract_city(osm_pbf_path)
-    sql = sql_tmpl.render(city=sql_safe(city), file_name=sql_safe(osm_pbf_path))
+    sql = sql_tmpl.render(city=sql_safe(city),
+                          file_name=sql_safe(osm_pbf_path),
+                          extract_date=str(extract_date))
     conn.execute(sql)
 
-def append_data_to_duckdb_file(duckdb_file: Path, osm_pbf_path: Path) -> None:
+def append_data_to_duckdb_file(duckdb_file: Path, osm_pbf_path: Path, extract_date: date) -> None:
     with duckdb.connect(str(duckdb_file)) as conn:
-        append_data_to_duckdb(conn, osm_pbf_path)
+        append_data_to_duckdb(conn, osm_pbf_path, extract_date)
 
 @click.command()
 @click.option('--output-file', type=click.Path(exists=False, file_okay=True, dir_okay=False, writable=True, path_type=Path))
 @click.option('--input-dir', type=click.Path(exists=False, file_okay=False, dir_okay=True, readable=True, path_type=Path))
-def create_duckdb_file(output_file: Path, input_dir: Path) -> None:
+@click.option('--date', '--date-end', type=click.DateTime(formats=["%Y-%m-%d"]))
+def create_duckdb_file(output_file: Path, input_dir: Path, date: datetime) -> None:
+
     output_file.parent.mkdir(parents=True, exist_ok=True)
     files = [f for f in input_dir.glob('**/*osm.pbf') if f.is_file()]
 
@@ -106,7 +118,7 @@ def create_duckdb_file(output_file: Path, input_dir: Path) -> None:
         if "'" in str(osm_file):
             continue
         try:
-            append_data_to_duckdb_file(output_file, osm_file)
+            append_data_to_duckdb_file(output_file, osm_file, date.date())
         except Exception as e:
             click.echo(f"Could not process {osm_file}")