Abstract json processing differences between duckdb and postgres

src/ldlite/_database/__init__.py

@@ -98,10 +98,11 @@ def record_history(self, history: LoadHistory) -> None: ...
 class TypedDatabase(Database, Generic[DB]):
     def __init__(self, conn_factory: Callable[[], DB]):
         self._conn_factory = conn_factory
-        with closing(self._conn_factory()) as conn, conn.cursor() as cur:
+        with closing(self._conn_factory()) as conn:
             try:
-                cur.execute('CREATE SCHEMA IF NOT EXISTS "ldlite_system";')
-                cur.execute("""
+                with conn.cursor() as cur:
+                    cur.execute('CREATE SCHEMA IF NOT EXISTS "ldlite_system";')
+                    cur.execute("""
 CREATE TABLE IF NOT EXISTS "ldlite_system"."load_history" (
     "table_name" TEXT UNIQUE
     ,"path" TEXT
@@ -112,6 +113,8 @@ def __init__(self, conn_factory: Callable[[], DB]):
     ,"index_complete_utc" TIMESTAMP
     ,"row_count" INTEGER
 );""")
+
+                self._setup_jfuncs(conn)
             except psycopg.errors.UniqueViolation:
                 # postgres throws this when multiple threads try to create
                 # the same resource even if CREATE IF NOT EXISTS was used
@@ -122,6 +125,10 @@ def __init__(self, conn_factory: Callable[[], DB]):
             else:
                 conn.commit()
 
+    @staticmethod
+    @abstractmethod
+    def _setup_jfuncs(conn: DB) -> None: ...
+
     @property
     @abstractmethod
     def _default_schema(self) -> str: ...

src/ldlite/_database/duckdb.py

@@ -8,6 +8,83 @@
 
 
 class DuckDbDatabase(TypedDatabase[duckdb.DuckDBPyConnection]):
+    @staticmethod
+    def _setup_jfuncs(conn: duckdb.DuckDBPyConnection) -> None:
+        with conn.cursor() as cur:
+            cur.execute("SELECT string_split(ltrim(version(),'v'), '.') AS has_lambda;")
+            if ver := cur.fetchone():
+                (ma, mi, _) = ver[0]
+                if int(ma) > 1 and int(mi) >= 3:
+                    cur.execute("SET lambda_syntax = 'ENABLE_SINGLE_ARROW';")
+
+        with conn.cursor() as cur:
+            cur.execute(
+                r"""
+CREATE OR REPLACE FUNCTION ldlite_system.jtype_of(j) AS
+    CASE coalesce(main.json_type(j), 'NULL')
+        WHEN 'VARCHAR' THEN 'string'
+        WHEN 'BIGINT' THEN 'number'
+        WHEN 'DOUBLE' THEN 'number'
+        WHEN 'UBIGINT' THEN 'number'
+        WHEN 'OBJECT' THEN 'object'
+        WHEN 'BOOLEAN' THEN 'boolean'
+        WHEN 'ARRAY' THEN 'array'
+        WHEN 'NULL' THEN 'null'
+        ELSE main.json_type(j)
+    END
+;
+
+CREATE OR REPLACE FUNCTION ldlite_system.jextract(j, p) AS
+    CASE ldlite_system.jtype_of(main.json_extract(j, p))
+        WHEN 'string' THEN
+            CASE
+                WHEN lower(main.json_extract_string(j, p)) = 'null' THEN 'null'::JSON
+                WHEN length(main.json_extract_string(j, p)) = 0 THEN 'null'::JSON
+                ELSE main.json_extract(j, p)
+            END
+        WHEN 'object' THEN
+            CASE
+                WHEN main.json_extract_string(j, p) = '{}' THEN 'null'::JSON
+                ELSE main.json_extract(j, p)
+            END
+        WHEN 'array' THEN
+            CASE
+                WHEN length(list_filter((main.json_extract(j, p))::JSON[], x -> x != 'null'::JSON)) = 0 THEN 'null'::JSON
+                ELSE list_filter((main.json_extract(j, p))::JSON[], x -> x != 'null'::JSON)
+            END
+        ELSE coalesce(main.json_extract(j, p), 'null'::JSON)
+    END
+;
+
+CREATE OR REPLACE FUNCTION ldlite_system.jextract_string(j, p) AS
+    main.json_extract_string(ldlite_system.jextract(j, p), '$')
+;
+
+CREATE OR REPLACE FUNCTION ldlite_system.jobject_keys(j) AS
+    unnest(main.json_keys(j))
+;
+
+CREATE OR REPLACE FUNCTION ldlite_system.jis_uuid(j) AS
+    CASE ldlite_system.jtype_of(j)
+        WHEN 'string' THEN regexp_full_match(main.json_extract_string(j, '$'), '^[a-fA-F0-9]{8}-[a-fA-F0-9]{4}-[1-5][a-fA-F0-9]{3}-[89abAB][a-fA-F0-9]{3}-[a-fA-F0-9]{12}$')
+        ELSE FALSE
+    END
+;
+
+CREATE OR REPLACE FUNCTION ldlite_system.jis_datetime(j) AS
+    CASE ldlite_system.jtype_of(j)
+        WHEN 'string' THEN regexp_full_match(main.json_extract_string(j, '$'), '^\d{4}-[01]\d-[0123]\dT[012]\d:[012345]\d:[012345]\d\.\d{3}(\+\d{2}:\d{2})?$')
+        ELSE FALSE
+    END
+;
+
+CREATE OR REPLACE FUNCTION ldlite_system.jis_float(j) AS
+    coalesce(main.json_type(j), 'NULL')='DOUBLE'
+;
+
+""",  # noqa: E501
+            )
+
     @property
     def _default_schema(self) -> str:
         return "main"

src/ldlite/_database/postgres.py

@@ -13,6 +13,83 @@ def __init__(self, dsn: str):
         # same sql between duckdb and postgres
         super().__init__(lambda: psycopg.connect(dsn, cursor_factory=psycopg.RawCursor))
 
+    @staticmethod
+    def _setup_jfuncs(conn: psycopg.Connection) -> None:
+        with conn.cursor() as cur:
+            cur.execute(
+                r"""
+CREATE OR REPLACE FUNCTION ldlite_system.jtype_of(j JSONB) RETURNS TEXT AS $$
+BEGIN
+    RETURN jsonb_typeof(j);
+END
+$$ LANGUAGE plpgsql;
+
+CREATE OR REPLACE FUNCTION ldlite_system.jextract(j JSONB, p TEXT) RETURNS JSONB AS $$
+BEGIN
+    RETURN CASE
+        WHEN ldlite_system.jtype_of(j->p) = 'string' THEN
+            CASE
+                WHEN lower(j->>p) = 'null' THEN 'null'::JSONB
+                WHEN length(j->>p) = 0 THEN 'null'::JSONB
+                ELSE j->p
+            END
+        WHEN ldlite_system.jtype_of(j->p) = 'array' THEN
+            CASE
+                WHEN jsonb_array_length(jsonb_path_query_array(j->p, '$[*] ? (@ != null)')) = 0 THEN 'null'::JSONB
+                ELSE jsonb_path_query_array(j->p, '$[*] ? (@ != null)')
+            END
+        WHEN ldlite_system.jtype_of(j->p) = 'object' THEN
+            CASE
+                WHEN j->>p = '{}' THEN 'null'::JSONB
+                ELSE j->p
+            END
+        ELSE j->p
+    END;
+END
+$$ LANGUAGE plpgsql;
+
+CREATE OR REPLACE FUNCTION ldlite_system.jextract_string(j JSONB, p TEXT) RETURNS TEXT AS $$
+BEGIN
+    RETURN ldlite_system.jextract(j, p) ->> 0;
+END
+$$ LANGUAGE plpgsql;
+
+CREATE OR REPLACE FUNCTION ldlite_system.jobject_keys(j JSONB) RETURNS SETOF TEXT AS $$
+BEGIN
+    RETURN QUERY SELECT jsonb_object_keys(j);
+END
+$$ LANGUAGE plpgsql;
+
+CREATE OR REPLACE FUNCTION ldlite_system.jis_uuid(j JSONB) RETURNS BOOLEAN AS $$
+BEGIN
+    RETURN CASE
+        WHEN ldlite_system.jtype_of(j) = 'string' THEN j->>0 ~ '^[a-fA-F0-9]{8}-[a-fA-F0-9]{4}-[1-5][a-fA-F0-9]{3}-[89abAB][a-fA-F0-9]{3}-[a-fA-F0-9]{12}$'
+        ELSE FALSE
+    END;
+END
+$$ LANGUAGE plpgsql;
+
+CREATE OR REPLACE FUNCTION ldlite_system.jis_datetime(j JSONB) RETURNS BOOLEAN AS $$
+BEGIN
+    RETURN CASE
+        WHEN ldlite_system.jtype_of(j) = 'string' THEN j->>0 ~ '^\d{4}-[01]\d-[0123]\dT[012]\d:[012345]\d:[012345]\d\.\d{3}(\+\d{2}:\d{2})?$'
+        ELSE FALSE
+    END;
+END
+$$ LANGUAGE plpgsql;
+
+
+CREATE OR REPLACE FUNCTION ldlite_system.jis_float(j JSONB) RETURNS BOOLEAN AS $$
+BEGIN
+    RETURN CASE
+        WHEN ldlite_system.jtype_of(j) = 'number' THEN j->>0 LIKE '%.%'
+        ELSE FALSE
+    END;
+END
+$$ LANGUAGE plpgsql;
+""",  # noqa: E501
+            )
+
     @property
     def _default_schema(self) -> str:
         return "public"