DoltTable Plugin

dev-requirements.txt

@@ -23,13 +23,11 @@ bcrypt==3.2.0
     # via
     #   -c requirements.txt
     #   paramiko
-black==20.8b1
+black==21.4b2
     # via
     #   -c requirements.txt
     #   -r dev-requirements.in
     #   flake8-black
-cached-property==1.5.2
-    # via docker-compose
 certifi==2020.12.5
     # via
     #   -c requirements.txt
@@ -59,7 +57,7 @@ cryptography==3.4.7
     # via
     #   -c requirements.txt
     #   paramiko
-dataclasses-json==0.5.2
+dataclasses-json==0.5.3
     # via
     #   -c requirements.txt
     #   flytekit
@@ -100,7 +98,7 @@ flake8==3.9.1
     #   -r dev-requirements.in
     #   flake8-black
     #   flake8-isort
-flyteidl==0.18.38
+flyteidl==0.18.39
     # via
     #   -c requirements.txt
     #   flytekit
@@ -115,11 +113,7 @@ idna==2.10
 importlib-metadata==4.0.1
     # via
     #   -c requirements.txt
-    #   flake8
-    #   jsonschema
     #   keyring
-    #   pluggy
-    #   pytest
 iniconfig==1.1.1
     # via pytest
 isort==5.8.0
@@ -240,7 +234,7 @@ python-dateutil==2.8.1
     #   croniter
     #   flytekit
     #   pandas
-python-dotenv==0.17.0
+python-dotenv==0.17.1
     # via docker-compose
 pytimeparse==1.1.8
     # via
@@ -267,7 +261,7 @@ requests==2.25.1
     #   docker-compose
     #   flytekit
     #   responses
-responses==0.13.2
+responses==0.13.3
     # via
     #   -c requirements.txt
     #   flytekit
@@ -316,15 +310,10 @@ toml==0.10.2
     #   coverage
     #   pytest
 typed-ast==1.4.3
-    # via
-    #   -c requirements.txt
-    #   black
-    #   mypy
+    # via mypy
 typing-extensions==3.7.4.3
     # via
     #   -c requirements.txt
-    #   black
-    #   importlib-metadata
     #   mypy
     #   typing-inspect
 typing-inspect==0.6.0

doc-requirements.txt

@@ -16,15 +16,15 @@ appnope==0.1.2
     # via
     #   ipykernel
     #   ipython
-astroid==2.5.3
+astroid==2.5.6
     # via sphinx-autoapi
 async-generator==1.10
     # via nbclient
 attrs==20.3.0
     # via
     #   jsonschema
     #   scantree
-babel==2.9.0
+babel==2.9.1
     # via sphinx
 backcall==0.2.0
     # via ipython
@@ -35,13 +35,13 @@ beautifulsoup4==4.9.3
     #   furo
     #   sphinx-code-include
     #   sphinx-material
-black==20.8b1
+black==21.4b2
     # via papermill
 bleach==3.3.0
     # via nbconvert
-boto3==1.17.55
+boto3==1.17.61
     # via sagemaker-training
-botocore==1.20.55
+botocore==1.20.61
     # via
     #   boto3
     #   s3transfer
@@ -68,7 +68,7 @@ cryptography==3.4.7
     #   paramiko
 css-html-js-minify==2.5.5
     # via sphinx-material
-dataclasses-json==0.5.2
+dataclasses-json==0.5.3
     # via flytekit
 decorator==5.0.7
     # via
@@ -88,7 +88,7 @@ entrypoints==0.3
     # via
     #   nbconvert
     #   papermill
-flyteidl==0.18.38
+flyteidl==0.18.39
     # via flytekit
 furo==2021.4.11b34
     # via -r doc-requirements.in
@@ -107,9 +107,7 @@ idna==2.10
 imagesize==1.2.0
     # via sphinx
 importlib-metadata==4.0.1
-    # via
-    #   jsonschema
-    #   keyring
+    # via keyring
 inotify_simple==1.2.1
     # via sagemaker-training
 ipykernel==5.5.3
@@ -279,19 +277,19 @@ requests==2.25.1
     #   papermill
     #   responses
     #   sphinx
-responses==0.13.2
+responses==0.13.3
     # via flytekit
 retry==0.9.2
     # via flytekit
 retrying==1.3.3
     # via sagemaker-training
-s3transfer==0.4.1
+s3transfer==0.4.2
     # via boto3
-sagemaker-training==3.9.1
+sagemaker-training==3.9.2
     # via flytekit
 scantree==0.0.1
     # via dirhash
-scipy==1.6.2
+scipy==1.6.3
     # via sagemaker-training
 six==1.15.0
     # via
@@ -316,13 +314,13 @@ sortedcontainers==2.3.0
     # via flytekit
 soupsieve==2.2.1
     # via beautifulsoup4
-sphinx-autoapi==1.8.0
+sphinx-autoapi==1.8.1
     # via -r doc-requirements.in
 sphinx-code-include==1.1.1
     # via -r doc-requirements.in
 sphinx-copybutton==0.3.1
     # via -r doc-requirements.in
-sphinx-gallery==0.8.2
+sphinx-gallery==0.9.0
     # via -r doc-requirements.in
 sphinx-material==0.0.32
     # via -r doc-requirements.in
@@ -381,15 +379,8 @@ traitlets==5.0.5
     #   nbclient
     #   nbconvert
     #   nbformat
-typed-ast==1.4.3
-    # via
-    #   astroid
-    #   black
 typing-extensions==3.7.4.3
-    # via
-    #   black
-    #   importlib-metadata
-    #   typing-inspect
+    # via typing-inspect
 typing-inspect==0.6.0
     # via dataclasses-json
 unidecode==1.2.0

flytekit/core/promise.py

@@ -621,7 +621,7 @@ def node_id(self):
 
     @property
     def node(self) -> Node:
-        """"""
+        """ """
         return self._node
 
     def __repr__(self) -> str:

plugins/flytekitplugins.dolt/flytekitplugins/dolt/__init__.py

@@ -0,0 +1 @@
+from .schema import DoltConfig, DoltTable, DoltTableNameTransformer

plugins/flytekitplugins.dolt/flytekitplugins/dolt/schema.py

@@ -0,0 +1,105 @@
+import tempfile
+import typing
+from dataclasses import dataclass
+from typing import Type
+
+import dolt_integrations.core as dolt_int
+import doltcli as dolt
+import pandas
+from dataclasses_json import dataclass_json
+from google.protobuf.struct_pb2 import Struct
+
+from flytekit import FlyteContext
+from flytekit.extend import TypeEngine, TypeTransformer
+from flytekit.models import types as _type_models
+from flytekit.models.literals import Literal, Scalar
+from flytekit.models.types import LiteralType
+
+
+@dataclass_json
+@dataclass
+class DoltConfig:
+    db_path: str
+    tablename: typing.Optional[str] = None
+    sql: typing.Optional[str] = None
+    io_args: typing.Optional[dict] = None
+    branch_conf: typing.Optional[dolt_int.Branch] = None
+    meta_conf: typing.Optional[dolt_int.Meta] = None
+    remote_conf: typing.Optional[dolt_int.Remote] = None
+
+
+@dataclass_json
+@dataclass
+class DoltTable:
+    config: DoltConfig
+    data: typing.Optional[pandas.DataFrame] = None
+
+
+class DoltTableNameTransformer(TypeTransformer[DoltTable]):
+    def __init__(self):
+        super().__init__(name="DoltTable", t=DoltTable)
+
+    def get_literal_type(self, t: Type[DoltTable]) -> LiteralType:
+        return LiteralType(simple=_type_models.SimpleType.STRUCT, metadata={})
+
+    def to_literal(
+        self,
+        ctx: FlyteContext,
+        python_val: DoltTable,
+        python_type: typing.Type[DoltTable],
+        expected: LiteralType,
+    ) -> Literal:
+
+        if not isinstance(python_val, DoltTable):
+            raise AssertionError(f"Value cannot be converted to a table: {python_val}")
+
+        conf = python_val.config
+        if python_val.data is not None and python_val.config.tablename is not None:
+            db = dolt.Dolt(conf.db_path)
+            with tempfile.NamedTemporaryFile() as f:
+                python_val.data.to_csv(f.name, index=False)
+                dolt_int.save(
+                    db=db,
+                    tablename=conf.tablename,
+                    filename=f.name,
+                    branch_conf=conf.branch_conf,
+                    meta_conf=conf.meta_conf,
+                    remote_conf=conf.remote_conf,
+                    save_args=conf.io_args,
+                )
+
+        s = Struct()
+        s.update(python_val.to_dict())
+        return Literal(Scalar(generic=s))
+
+    def to_python_value(
+        self,
+        ctx: FlyteContext,
+        lv: Literal,
+        expected_python_type: typing.Type[DoltTable],
+    ) -> DoltTable:
+
+        if not (lv and lv.scalar and lv.scalar.generic and lv.scalar.generic["config"]):
+            return pandas.DataFrame()
+
+        conf = DoltConfig(**lv.scalar.generic["config"])
+        db = dolt.Dolt(conf.db_path)
+
+        with tempfile.NamedTemporaryFile() as f:
+            dolt_int.load(
+                db=db,
+                tablename=conf.tablename,
+                sql=conf.sql,
+                filename=f.name,
+                branch_conf=conf.branch_conf,
+                meta_conf=conf.meta_conf,
+                remote_conf=conf.remote_conf,
+                load_args=conf.io_args,
+            )
+            df = pandas.read_csv(f)
+            lv.data = df
+
+        return lv
+
+
+TypeEngine.register(DoltTableNameTransformer())

plugins/flytekitplugins.dolt/scripts/flytekit_install_dolt.sh

@@ -0,0 +1,12 @@
+#!/bin/bash
+
+# Fetches and install Dolt. To be invoked by the Dockerfile
+
+# echo commands to the terminal output
+set -eox pipefail
+
+# Install Dolt
+
+apt-get update -y \
+    && apt-get install curl \
+    && sudo bash -c 'curl -L https://github.com/dolthub/dolt/releases/latest/download/install.sh | sudo bash'

plugins/flytekitplugins.dolt/setup.py

@@ -0,0 +1,66 @@
+import shlex
+import subprocess
+import urllib.request
+
+from setuptools import setup
+from setuptools.command.develop import develop
+
+PLUGIN_NAME = "dolt"
+
+microlib_name = f"flytekitplugins.{PLUGIN_NAME}"
+
+plugin_requires = ["flytekit>=0.16.0b0,<1.0.0", "dolt_integrations>=0.1.3"]
+
+__version__ = "0.0.0+develop"
+
+
+class PostDevelopCommand(develop):
+    """Post-installation for development mode."""
+
+    def run(self):
+        develop.run(self)
+        install, _ = urllib.request.urlretrieve(
+            "https://github.com/liquidata-inc/dolt/releases/latest/download/install.sh"
+        )
+        subprocess.call(shlex.split(f"chmod +x {install}"))
+        subprocess.call(shlex.split(f"sudo {install}"))
+
+        pref = "dolt config --global --add"
+        subprocess.call(
+            shlex.split(f"{pref} user.email [email protected]"),
+        )
+        subprocess.call(
+            shlex.split(f"{pref} user.name 'Bojack Horseman'"),
+        )
+        subprocess.call(
+            shlex.split(f"{pref} metrics.host eventsapi.awsdev.ld-corp.com"),
+        )
+        subprocess.call(shlex.split(f"{pref} metrics.port 443"))
+
+
+setup(
+    name=microlib_name,
+    version=__version__,
+    author="dolthub",
+    author_email="[email protected]",
+    description="Dolt plugin for flytekit",
+    namespace_packages=["flytekitplugins"],
+    packages=[f"flytekitplugins.{PLUGIN_NAME}"],
+    install_requires=plugin_requires,
+    cmdclass=dict(develop=PostDevelopCommand),
+    license="apache2",
+    python_requires=">=3.7",
+    classifiers=[
+        "Intended Audience :: Science/Research",
+        "Intended Audience :: Developers",
+        "License :: OSI Approved :: Apache Software License",
+        "Programming Language :: Python :: 3.7",
+        "Programming Language :: Python :: 3.8",
+        "Topic :: Scientific/Engineering",
+        "Topic :: Scientific/Engineering :: Artificial Intelligence",
+        "Topic :: Software Development",
+        "Topic :: Software Development :: Libraries",
+        "Topic :: Software Development :: Libraries :: Python Modules",
+    ],
+    scripts=["scripts/flytekit_install_dolt.sh"],
+)