chg: use Duckdb Python API for GBIF backbone

Author: frafra

src/datasync/gbif_backbone.py

@@ -1,20 +1,16 @@
-from importlib.resources import files
-
-import duckdb
 import fsspec
 import typer
+from duckdb import (
+    DuckDBPyConnection,
+    connect,
+)
 
-from .libs.helpers import pa_read_tsv
+from .libs.helpers import DuckDBAtomicTransaction
 from .settings import (
     env,
     log,
 )
 
-resources = files(__package__).joinpath("gbif_backbone")
-taxon_sql = resources.joinpath("taxon.sql").read_text()
-vernacularname_sql = resources.joinpath("vernacularname.sql").read_text()
-
-
 log.debug("Importing GBIF Backbone settings")
 
 
@@ -29,23 +25,41 @@
 app = typer.Typer(help="export GBIF Backbone data to DuckDB database")
 
 
-def import_taxon(conn, archive):
+def import_taxon(conn: DuckDBPyConnection, archive):
     log.debug("Importing Taxon.tsv")
-    pa_taxon = pa_read_tsv(archive, "Taxon.tsv")  # noqa: F841
-    conn.execute(taxon_sql)
+    conn.sql("DROP TABLE IF EXISTS taxon")
+    conn.from_csv_auto(archive.open("Taxon.tsv")).to_table("taxon")
+    conn.sql("CREATE INDEX taxon_taxonid ON taxon (taxonid)")
+    conn.sql("""
+        PRAGMA create_fts_index(
+            "taxon", "taxonID", "canonicalName", overwrite=TRUE
+        )
+    """)
 
 
-def import_vernacular_name(conn, archive):
+def import_vernacular_name(conn: DuckDBPyConnection, archive):
     log.debug("Importing VernacularName.tsv")
-    pa_vernacular_names = pa_read_tsv(archive, "VernacularName.tsv")  # noqa: F841
-    conn.execute(vernacularname_sql)
+    conn.sql("DROP TABLE IF EXISTS vernacular_name")
+    vernacular_names = conn.from_csv_auto(archive.open("VernacularName.tsv"))  # noqa: F841
+    conn.sql("""
+        SELECT *, CONCAT_WS('|', taxonID, language, vernacularName) AS vernacularID
+        FROM vernacular_names
+    """).to_table("vernacular_name")
+    conn.sql("CREATE INDEX vernacular_name_taxonid ON vernacular_name (taxonid)")
+    conn.sql("CREATE INDEX vernacular_name_language ON vernacular_name (language)")
+    conn.sql("""
+        PRAGMA create_fts_index(
+            "vernacular_name", "vernacularID", "vernacularName", overwrite=TRUE
+        )
+    """)
 
 
 @app.command()
 def import_all():
     """Import GBIF Backbone data into a DuckDB database."""
     archive = fsspec.filesystem("zip", fo=GBIF_BACKBONE_URL.geturl(), mode="r")
-    with duckdb.connect(GBIF_BACKBONE_DUCKDB_NAME) as conn:
-        import_taxon(conn, archive)
-        import_vernacular_name(conn, archive)
+    with connect(GBIF_BACKBONE_DUCKDB_NAME) as conn:
+        with DuckDBAtomicTransaction(conn):
+            import_taxon(conn, archive)
+            import_vernacular_name(conn, archive)
     log.info("GBIF Backbone data imported successfully")

src/datasync/gbif_backbone/taxon.sql

@@ -1,8 +1,8 @@
 import re
 
+from duckdb import DuckDBPyConnection
 from lxml import etree
 from lxml.etree import _Element
-from pyarrow import csv
 
 PARSER: etree.XMLParser = etree.XMLParser(resolve_entities=False)
 
@@ -32,13 +32,16 @@ def get_anytext(bag: str | _Element | list[str]) -> str:
         raise TypeError("xpath result was not a list of strings")
 
 
-TSV_PARSE_OPTIONS = csv.ParseOptions(
-    delimiter="\t", quote_char=False, double_quote=False
-)
+class DuckDBAtomicTransaction:
+    def __init__(self, conn: DuckDBPyConnection):
+        self.conn = conn
 
+    def __enter__(self):
+        self.conn.begin()
+        return self.conn
 
-def pa_read_tsv(archive, filename):
-    return csv.read_csv(
-        archive.open(filename),
-        parse_options=TSV_PARSE_OPTIONS,
-    )
+    def __exit__(self, exc_type, exc_value, traceback):
+        if exc_type is None:
+            self.conn.commit()
+        else:
+            self.conn.rollback()