db_sync.py

import collections.abc
import csv
import itertools
import json
import os
import re
import time
import uuid
from urllib.parse import urlparse

from target_duckdb.logger import get_logger


# copied from inflection.camelize to eliminate a dependency
def camelize(string: str, uppercase_first_letter: bool = True) -> str:
    if uppercase_first_letter:
        return re.sub(r"(?:^|_)(.)", lambda m: m.group(1).upper(), string)
    else:
        return string[0].lower() + camelize(string)[1:]


# pylint: disable=fixme
def column_type(schema_property):
    property_type = schema_property["type"]
    property_format = schema_property["format"] if "format" in schema_property else None
    col_type = "varchar"
    if "object" in property_type or "array" in property_type:
        col_type = "json"

    # Every date-time JSON value is currently mapped to TIMESTAMP (no time zone)
    elif property_format == "date-time":
        col_type = "timestamp"
    elif property_format == "time":
        col_type = "time"
    elif "number" in property_type:
        col_type = "double"
    elif "integer" in property_type and "string" in property_type:
        col_type = "varchar"
    elif "integer" in property_type:
        if "maximum" in schema_property:
            if schema_property["maximum"] <= 32767:
                col_type = "smallint"
            elif schema_property["maximum"] <= 2147483647:
                col_type = "integer"
            elif schema_property["maximum"] <= 9223372036854775807:
                col_type = "bigint"
        else:
            col_type = "hugeint"
    elif "boolean" in property_type:
        col_type = "boolean"

    get_logger("target_duckdb").debug(
        "schema_property: %s -> col_type: %s", schema_property, col_type
    )

    return col_type


def safe_column_name(name):
    return '"{}"'.format(name).lower()


def column_clause(name, schema_property):
    return "{} {}".format(safe_column_name(name), column_type(schema_property))


def flatten_key(k, parent_key, sep):
    return sep.join(parent_key + [k])


# pylint: disable=dangerous-default-value,invalid-name
def flatten_schema(d, parent_key=[], sep="__", level=0, max_level=0):
    items = []

    if "properties" not in d:
        return {}

    for k, v in d["properties"].items():
        new_key = flatten_key(k, parent_key, sep)
        if "type" in v.keys():
            if "object" in v["type"] and "properties" in v and level < max_level:
                items.extend(
                    flatten_schema(
                        v,
                        parent_key + [k],
                        sep=sep,
                        level=level + 1,
                        max_level=max_level,
                    ).items()
                )
            else:
                items.append((new_key, v))
        elif "anyOf" in v.keys():
            llv = v["anyOf"]
            i, entry = 0, llv[0]
            while entry["type"] == "null":
                i += 1
                entry = llv[i]
            entry["type"] = ["null", entry["type"]]
            items.append((new_key, entry))
        else:
            if len(v.values()) > 0:
                if list(v.values())[0][0]["type"] == "string":
                    list(v.values())[0][0]["type"] = ["null", "string"]
                    items.append((new_key, list(v.values())[0][0]))
                elif list(v.values())[0][0]["type"] == "array":
                    list(v.values())[0][0]["type"] = ["null", "array"]
                    items.append((new_key, list(v.values())[0][0]))
                elif list(v.values())[0][0]["type"] == "object":
                    list(v.values())[0][0]["type"] = ["null", "object"]
                    items.append((new_key, list(v.values())[0][0]))

    key_func = lambda item: item[0]
    sorted_items = sorted(items, key=key_func)
    for k, g in itertools.groupby(sorted_items, key=key_func):
        if len(list(g)) > 1:
            raise ValueError("Duplicate column name produced in schema: {}".format(k))

    return dict(sorted_items)


# pylint: disable=redefined-outer-name
def _should_json_dump_value(key, value, flatten_schema=None):
    if isinstance(value, (dict, list)):
        return True

    if (
        flatten_schema
        and key in flatten_schema
        and "type" in flatten_schema[key]
        and set(flatten_schema[key]["type"]) == {"null", "object", "array"}
    ):
        return True

    return False


# pylint: disable-msg=too-many-arguments
def flatten_record(
    d, flatten_schema=None, parent_key=[], sep="__", level=0, max_level=0
):
    items = []
    for k, v in d.items():
        new_key = flatten_key(k, parent_key, sep)
        if isinstance(v, collections.abc.MutableMapping) and level < max_level:
            items.extend(
                flatten_record(
                    v,
                    flatten_schema,
                    parent_key + [k],
                    sep=sep,
                    level=level + 1,
                    max_level=max_level,
                ).items()
            )
        else:
            items.append(
                (
                    new_key,
                    json.dumps(v)
                    if _should_json_dump_value(k, v, flatten_schema)
                    else v,
                )
            )
    return dict(items)


def primary_column_names(stream_schema_message):
    return [safe_column_name(p) for p in stream_schema_message["key_properties"]]


def stream_name_to_dict(stream_name, separator="-"):
    catalog_name = None
    schema_name = None
    table_name = stream_name

    # Schema and table name can be derived from stream if it's in <schema_nama>-<table_name> format
    s = stream_name.split(separator)
    if len(s) == 2:
        schema_name = s[0]
        table_name = s[1]
    if len(s) > 2:
        catalog_name = s[0]
        schema_name = s[1]
        table_name = "_".join(s[2:])

    return {
        "catalog_name": catalog_name,
        "schema_name": schema_name,
        "table_name": table_name,
    }


def get_catalog_name(connection_config):
    raw_path = connection_config.get("path", connection_config.get("filepath"))
    parsed = urlparse(raw_path)
    base_file = os.path.basename(parsed.path)
    path_db = os.path.splitext(base_file)[0]
    if parsed.scheme in {"md", "motherduck"} and path_db == "":
        path_db = "my_db"
    return path_db


# pylint: disable=too-many-public-methods,too-many-instance-attributes
class DbSync:
    def __init__(self, connection, connection_config, stream_schema_message=None):
        """
        connection:  DuckDB connection
        connection_config: Connection config information
        stream_schema_message:  An instance of the DbSync class is typically used to load
                                data only from a certain singer tap stream.

                                The stream_schema_message holds the destination schema
                                name and the JSON schema that will be used to
                                validate every RECORDS messages that comes from the stream.
                                Schema validation happening before creating CSV and before
                                loading data into DuckDB.

                                If stream_schema_message is not defined then we can use
                                the DbSync instance as a generic purpose connection to
                                DuckDB and can run individual queries. For example
                                collecting catalog information from DuckDB for caching
                                purposes.
        """
        self.conn = connection
        self.connection_config = connection_config
        self.stream_schema_message = stream_schema_message

        # logger to be used across the class's methods
        self.logger = get_logger("target_duckdb")

        # setup a catalog name
        if self.connection_config.get("database"):
            self.catalog_name = self.connection_config.get("database")
        elif self.connection_config.get("dbname"):
            self.catalog_name = self.connection_config.get("dbname")
        else:
            self.catalog_name = get_catalog_name(self.connection_config)
        self.schema_name = None

        # Init stream schema
        if stream_schema_message is not None:
            # Define initial list of indices to created
            self.hard_delete = self.connection_config.get("hard_delete")
            if self.hard_delete:
                self.indices = ["_sdc_deleted_at"]
            else:
                self.indices = []

            #  Define target schema name.
            #  --------------------------
            #  Target schema name can be defined in multiple ways:
            #
            #   1: 'default_target_schema' key  : Target schema is the same for every incoming stream if
            #       not specified explicitly for a given stream in the `schema_mapping` object
            #   2: 'schema_mapping' key : Target schema defined explicitly for a given stream.
            #       Example config.json:
            #           "schema_mapping": {
            #               "my_tap_stream_id": {
            #                   "target_schema": "my_duckdb_schema",
            #                   "indices": ["column_1", "column_2s"]
            #               }
            #           }
            config_default_target_schema = self.connection_config.get(
                "default_target_schema", ""
            ).strip()
            config_schema_mapping = self.connection_config.get("schema_mapping", {})

            stream_name = stream_schema_message["stream"]
            stream_schema_name = stream_name_to_dict(stream_name)["schema_name"]
            stream_table_name = stream_name_to_dict(stream_name)["table_name"]
            if config_schema_mapping and stream_schema_name in config_schema_mapping:
                self.schema_name = config_schema_mapping[stream_schema_name].get(
                    "target_schema"
                )

                # Get indices to create for the target table
                indices = config_schema_mapping[stream_schema_name].get("indices", {})
                if stream_table_name in indices:
                    self.indices.extend(indices.get(stream_table_name, []))

            elif config_default_target_schema:
                self.schema_name = config_default_target_schema

            if not self.schema_name:
                raise Exception(
                    "Target schema name not defined in config. Neither 'default_target_schema' (string)"
                    "nor 'schema_mapping' (object) defines target schema for {} stream.".format(
                        stream_name
                    )
                )

            self.data_flattening_max_level = self.connection_config.get(
                "data_flattening_max_level", 0
            )
            self.flatten_schema = flatten_schema(
                stream_schema_message["schema"],
                max_level=self.data_flattening_max_level,
            )

            # Delimiters/quotechars for CSV temp files
            self.delimiter = self.connection_config.get("delimiter", ",")
            self.quotechar = self.connection_config.get("quotechar", '"')

    def query(self, query, params=None):
        self.logger.debug("Running query: %s", query)
        cur = self.conn
        if params:
            cur.execute(query, params)
        else:
            cur.execute(query)

        cols = [x[0] for x in cur.description]
        ret = []
        for row in cur.fetchall():
            ret.append({cols[i]: row[i] for i in range(len(cols))})
        return ret

    def table_name(self, stream_name, is_temporary=False, without_schema=False):
        if is_temporary:
            return "tmp_{}".format(str(uuid.uuid4()).replace("-", "_"))

        stream_dict = stream_name_to_dict(stream_name)
        table_name = stream_dict["table_name"]
        pg_table_name = table_name.replace(".", "_").replace("-", "_").lower()

        if without_schema:
            return f'"{pg_table_name.lower()}"'
        elif self.catalog_name:
            return f'"{self.catalog_name}"."{self.schema_name}"."{pg_table_name.lower()}"'
        else:
            return f'"{self.schema_name}"."{pg_table_name.lower()}"'

    def record_primary_key_string(self, record):
        if len(self.stream_schema_message["key_properties"]) == 0:
            return None
        flatten = flatten_record(
            record, self.flatten_schema, max_level=self.data_flattening_max_level
        )
        try:
            key_props = [
                str(flatten[p]) for p in self.stream_schema_message["key_properties"]
            ]
        except Exception as exc:
            self.logger.info(
                "Cannot find %s primary key(s) in record: %s",
                self.stream_schema_message["key_properties"],
                flatten,
            )
            raise exc
        return ",".join(key_props)

    def record_to_flattened(self, record):
        flatten = flatten_record(
            record, self.flatten_schema, max_level=self.data_flattening_max_level
        )

        return [
            flatten[name]
            if name in flatten and (flatten[name] == 0 or flatten[name])
            else None
            for name in self.flatten_schema
        ]

    def load_rows(self, records, count, temp_dir):
        stream_schema_message = self.stream_schema_message
        stream = stream_schema_message["stream"]
        self.logger.info(
            "Loading %d rows into '%s'", count, self.table_name(stream, False)
        )

        cur = self.conn
        temp_table = self.table_name(stream_schema_message["stream"], is_temporary=True)
        temp_file_csv = os.path.join(temp_dir, f"{temp_table}.csv")
        cur.execute(self.create_table_query(table_name=temp_table, is_temporary=True))

        # batch the records into a CSV file and do a copy operation
        with open(temp_file_csv, "w") as f:
            csvwriter = csv.writer(
                f,
                delimiter=self.delimiter,
                quotechar=self.quotechar,
                quoting=csv.QUOTE_MINIMAL,
            )
            for record in records:
                csvwriter.writerow(self.record_to_flattened(record))
        cur.execute("COPY {} FROM '{}' WITH (new_line '\\r\\n')".format(temp_table, temp_file_csv))

        if len(self.stream_schema_message["key_properties"]) > 0:
            cur.execute(self.update_from_temp_table(temp_table))
        cur.execute(self.insert_from_temp_table(temp_table))
        cur.execute(f"DROP TABLE {temp_table}")
        os.unlink(temp_file_csv)

    # pylint: disable=duplicate-string-formatting-argument
    def insert_from_temp_table(self, temp_table):
        stream_schema_message = self.stream_schema_message
        columns = self.column_names()
        table = self.table_name(stream_schema_message["stream"])

        if len(stream_schema_message["key_properties"]) == 0:
            return """INSERT INTO {} ({})
                    (SELECT s.* FROM {} s)
                    """.format(
                table, ", ".join(columns), temp_table
            )

        return """INSERT INTO {} ({})
        (SELECT s.* FROM {} s LEFT OUTER JOIN {} t ON {} WHERE {})
        """.format(
            table,
            ", ".join(columns),
            temp_table,
            table,
            self.primary_key_condition("t"),
            self.primary_key_null_condition("t"),
        )

    def update_from_temp_table(self, temp_table):
        stream_schema_message = self.stream_schema_message
        primary_key_columns = set(primary_column_names(self.stream_schema_message))
        columns = [x for x in self.column_names() if x not in primary_key_columns]
        table = self.table_name(stream_schema_message["stream"])

        if not columns:
            # Nothing to do if all of the columns are PKs
            return

        return """UPDATE {} as s SET {} FROM {} t
        WHERE {}
        """.format(
            table,
            ", ".join(["{}=t.{}".format(c, c) for c in columns]),
            temp_table,
            self.primary_key_condition("t"),
        )

    def primary_key_condition(self, right_table):
        stream_schema_message = self.stream_schema_message
        names = primary_column_names(stream_schema_message)
        return " AND ".join(["s.{} = {}.{}".format(c, right_table, c) for c in names])

    def primary_key_null_condition(self, right_table):
        stream_schema_message = self.stream_schema_message
        names = primary_column_names(stream_schema_message)
        return " AND ".join(["{}.{} is null".format(right_table, c) for c in names])

    def column_names(self):
        return [safe_column_name(name) for name in self.flatten_schema]

    def column_types(self):
        return [column_type(schema) for (name, schema) in self.flatten_schema.items()]

    def create_table_query(self, table_name=None, is_temporary=False):
        stream_schema_message = self.stream_schema_message
        columns = [
            column_clause(name, schema)
            for (name, schema) in self.flatten_schema.items()
        ]

        if len(stream_schema_message["key_properties"]) > 0:
            primary_key = [
                "PRIMARY KEY ({})".format(
                    ", ".join(primary_column_names(stream_schema_message))
                )
            ]
        else:
            primary_key = []

        if not table_name:
            gen_table_name = self.table_name(
                stream_schema_message["stream"], is_temporary=is_temporary
            )

        return "CREATE TABLE IF NOT EXISTS {} ({})".format(
            table_name if table_name else gen_table_name,
            ", ".join(columns + primary_key),
        )

    def create_index(self, stream, column):
        table = self.table_name(stream)
        table_without_schema = self.table_name(stream, without_schema=True)
        index_name = "i_{}_{}".format(
            table_without_schema[:30].replace(" ", "").replace('"', ""),
            column.replace(",", "_"),
        )
        query = "CREATE INDEX IF NOT EXISTS {} ON {} ({})".format(
            index_name, table, column
        )
        self.logger.info(
            "Creating index on '%s' table on '%s' column(s)... %s", table, column, query
        )
        self.query(query)

    def create_indices(self, stream):
        if isinstance(self.indices, list):
            for index in self.indices:
                self.create_index(stream, index)

    def delete_rows(self, stream):
        table = self.table_name(stream)
        query = "DELETE FROM {} WHERE _sdc_deleted_at IS NOT NULL RETURNING _sdc_deleted_at".format(
            table
        )
        self.logger.info("Deleting rows from '%s' table... %s", table, query)
        self.logger.info("DELETE %s", len(self.query(query)))

    def create_schema_if_not_exists(self, table_columns_cache=None):
        schema_name = self.schema_name
        catalog_name = self.catalog_name
        schema_rows = 0

        # table_columns_cache is an optional pre-collected list of available objects in DuckDB
        if table_columns_cache:
            schema_rows = list(
                filter(lambda x: x["TABLE_SCHEMA"] == schema_name, table_columns_cache)
            )
        # Query realtime if not pre-collected
        elif catalog_name:
            schema_rows = self.query(
                "SELECT LOWER(schema_name) schema_name, LOWER(catalog_name) catalog_name "
                "FROM information_schema.schemata "
                "WHERE LOWER(schema_name) = ? AND LOWER(catalog_name) = ?",
                (
                    schema_name.lower(),
                    catalog_name.lower(),
                ),
            )
        else:
            schema_rows = self.query(
                "SELECT LOWER(schema_name) schema_name FROM information_schema.schemata WHERE LOWER(schema_name) = ?",
                (schema_name.lower(),),
            )

        if len(schema_rows) == 0:
            if catalog_name:
                query = f'CREATE SCHEMA IF NOT EXISTS "{catalog_name}"."{schema_name}"'
            else:
                query = f'CREATE SCHEMA IF NOT EXISTS "{schema_name}"'

            self.logger.info(
                "Schema '%s' does not exist. Creating... %s", schema_name, query
            )
            self.query(query)

    def get_tables(self):
        if self.catalog_name:
            return self.query(
                "SELECT table_name FROM information_schema.tables WHERE table_catalog = ? AND table_schema = ?",
                (self.catalog_name, self.schema_name),
            )
        else:
            return self.query(
                "SELECT table_name FROM information_schema.tables WHERE table_schema = ?",
                (self.schema_name,),
            )

    def get_table_columns(self, table_name):
        return self.query(
            """SELECT column_name, data_type
        FROM information_schema.columns
        WHERE lower(table_name) = ? AND lower(table_schema) = ?""",
            (table_name.replace('"', "").lower(), self.schema_name.lower()),
        )

    def update_columns(self):
        stream_schema_message = self.stream_schema_message
        stream = stream_schema_message["stream"]
        table_name = self.table_name(stream, without_schema=True)
        columns = self.get_table_columns(table_name)
        columns_dict = {column["column_name"].lower(): column for column in columns}

        columns_to_add = [
            column_clause(name, properties_schema)
            for (name, properties_schema) in self.flatten_schema.items()
            if name.lower() not in columns_dict
        ]

        for column in columns_to_add:
            self.add_column(column, stream)

        columns_to_replace = [
            (safe_column_name(name), column_clause(name, properties_schema))
            for (name, properties_schema) in self.flatten_schema.items()
            if name.lower() in columns_dict
            and columns_dict[name.lower()]["data_type"].lower()
            != column_type(properties_schema).lower()
        ]

        for column_name, column in columns_to_replace:
            self.version_column(column_name, stream)
            self.add_column(column, stream)

    def drop_column(self, column_name, stream):
        drop_column = "ALTER TABLE {} DROP COLUMN {}".format(
            self.table_name(stream), column_name
        )
        self.logger.info("Dropping column: %s", drop_column)
        self.query(drop_column)

    def version_column(self, column_name, stream):
        version_column = 'ALTER TABLE {} RENAME COLUMN {} TO "{}_{}"'.format(
            self.table_name(stream, False),
            column_name,
            column_name.replace('"', ""),
            time.strftime("%Y%m%d_%H%M"),
        )
        self.logger.info("Versioning column: %s", version_column)
        self.query(version_column)

    def add_column(self, column, stream):
        add_column = "ALTER TABLE {} ADD COLUMN {}".format(
            self.table_name(stream), column
        )
        self.logger.info("Adding column: %s", add_column)
        self.query(add_column)

    def sync_table(self):
        stream_schema_message = self.stream_schema_message
        stream = stream_schema_message["stream"]
        table_name = self.table_name(stream, without_schema=True)
        found_tables = [
            table
            for table in (self.get_tables())
            if f'"{table["table_name"].lower()}"' == table_name
        ]
        if len(found_tables) == 0:
            query = self.create_table_query()
            self.logger.info(
                "Table '%s' does not exist. Creating... %s", table_name, query
            )
            self.query(query)
        else:
            self.logger.info("Table '%s' exists", table_name)
            self.update_columns()