probe3.py

#!/usr/bin/env python3
"""probe3.py - exercise the exotic SQL types.

probe2.py covered Bool/SmallInt/Integer/BigInt/Date/Timestamp/Double/Varchar.
This adds everything else the Hyper API can construct: Char, Varchar(n), Bytes,
Json, Geography, Interval, Time, TimestampTZ, Numeric at three precisions, Oid
and Float(32).

One file per type so a type the API refuses to insert does not take the rest
down with it. Each file has 6 columns - cardinality {3, 20, 300} x {no nulls,
nulls} - which between them reach the 2-bit, 8-bit and 16-bit code widths in
both the non-null and nullable scheme families.
"""
import datetime
import decimal
import os
import shutil
import sys

from tableauhyperapi import (HyperProcess, Connection, Telemetry, CreateMode,
                             TableDefinition, SqlType, Inserter, NULLABLE,
                             TableName, Interval)

HERE = os.path.dirname(os.path.abspath(__file__))
OUT = os.path.join(HERE, "probes3")
N = 1000
CARDS = [3, 20, 300]

# label -> (SqlType, value factory over the distinct index)
TYPES = {
    "char":   (SqlType.char(8), lambda k: f"c{k:07d}"),
    "vchar":  (SqlType.varchar(32), lambda k: f"v{k:04d}"),
    "bytes":  (SqlType.bytes(), lambda k: bytes([k & 0xFF, (k >> 8) & 0xFF, 0xAB, 0xCD])),
    "json":   (SqlType.json(), lambda k: f'{{"k":{k}}}'),
    "geo":    (SqlType.geography(), None),          # inserted via CAST, see below
    "intvl":  (SqlType.interval(), lambda k: Interval(0, k, 0)),
    "time":   (SqlType.time(), lambda k: datetime.time(k // 3600 % 24,
                                                       k // 60 % 60, k % 60)),
    "tstz":   (SqlType.timestamp_tz(), lambda k: datetime.datetime(
        2020, 1, 1, tzinfo=datetime.timezone.utc) + datetime.timedelta(hours=k)),
    "num9":   (SqlType.numeric(9, 2), lambda k: decimal.Decimal(f"{1000 + k}.25")),
    "num18":  (SqlType.numeric(18, 4), lambda k: decimal.Decimal(f"{10**12 + k}.0625")),
    "num38":  (SqlType.numeric(38, 6), lambda k: decimal.Decimal(f"{10**24 + k}.015625")),
    "oid":    (SqlType.oid(), lambda k: 70000 + k),
    "float":  (SqlType.float(), lambda k: 1000.5 + k),
    # Negative values: every integer type is two's complement, and an unsigned
    # read turns -1 into 18446744073709551615. Nothing in the sample corpus is
    # negative, so only a probe can catch it.
    "negbig": (SqlType.big_int(), lambda k: -(10 ** 12) - k),
    "negsml": (SqlType.small_int(), lambda k: -30000 + k),
    "negnum": (SqlType.numeric(18, 4), lambda k: decimal.Decimal(f"-{10**9 + k}.5")),
}


def build(hp, label, sqltype, make):
    path = os.path.join(OUT, f"{label}.hyper")
    cols = []
    for card in CARDS:
        for nulls in (False, True):
            name = f"c{card}_{'n' if nulls else 'x'}"
            vals = [make(i % card) for i in range(N)]
            if nulls:
                for i in range(0, N, 7):
                    vals[i] = None
            cols.append((name, vals))
    tdef = TableDefinition(TableName("Extract", "E"),
                           [TableDefinition.Column(n, sqltype, NULLABLE)
                            for n, _ in cols])
    with Connection(hp.endpoint, path, CreateMode.CREATE_AND_REPLACE) as c:
        c.catalog.create_schema("Extract")
        c.catalog.create_table(tdef)
        with Inserter(c, tdef) as ins:
            ins.add_rows([[v[i] for _, v in cols] for i in range(N)])
            ins.execute()
    return path


def build_geography(hp):
    """Geography has no Python binding for insertion; go through SQL text."""
    path = os.path.join(OUT, "geo.hyper")
    with Connection(hp.endpoint, path, CreateMode.CREATE_AND_REPLACE) as c:
        c.execute_command('CREATE SCHEMA "Extract"')
        c.execute_command('CREATE TABLE "Extract"."E" '
                          '("c3_x" GEOGRAPHY, "c3_n" GEOGRAPHY)')
        rows = []
        for i in range(N):
            a = f"CAST('point({i % 3} {i % 3})' AS GEOGRAPHY)"
            b = "NULL" if i % 7 == 0 else f"CAST('point({i % 3} 1)' AS GEOGRAPHY)"
            rows.append(f"({a}, {b})")
        vals = ", ".join(rows)
        c.execute_command(f'INSERT INTO "Extract"."E" VALUES {vals}')
    return path


def main():
    shutil.rmtree(OUT, ignore_errors=True)
    os.makedirs(OUT)
    made, failed = [], []
    # numeric128 needs format v3 and float32 needs v4 (see README section 5),
    # and the default for creation is lower - so ask for the newest we know of.
    params = {"default_database_version": "4"}
    with HyperProcess(telemetry=Telemetry.DO_NOT_SEND_USAGE_DATA_TO_TABLEAU,
                      parameters=params) as hp:
        for label, (sqltype, make) in TYPES.items():
            try:
                p = build_geography(hp) if make is None else build(hp, label, sqltype, make)
                made.append(os.path.basename(p))
            except Exception as e:                    # keep going across types
                failed.append(f"{label}: {type(e).__name__}: {str(e).splitlines()[0][:90]}")
    for m in made:
        print(f"  ok   {m}")
    for m in failed:
        print(f"  FAIL {m}")
    print(f"\n{len(made)} built, {len(failed)} failed -> {OUT}")
    return 0


if __name__ == "__main__":
    sys.exit(main())