01_standalone_examples/azure-databricks/Unity Catalog Integration Demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "d773767c-565a-415a-9cee-7feecc2d5e68",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "# [Integration of lakeFS with Unity Catalog](https://docs.lakefs.io/integrations/unity-catalog.html)\n",
    "\n",
    "## Use Case: Isolated Development and Testing Environment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "24a903fa-02e3-4dac-9407-749e9c6bf2de",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "# Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "b5064bb5-fcc3-44af-895e-9c96cd338720",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "#### Run demo [Setup](./?o=8911673420610391#notebook/1786061902497011) tasks here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "3b243b52-2cc0-4820-849c-c7d1f7f60e6d",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "%run ./unityCatalogIntegrationDemoSetup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "1b98bce2-8c09-4b50-a6b6-dac0b6bd4dde",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "#### Let’s define the table descriptor and upload it to lakeFS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "7dec63d4-5403-439f-9959-de9754399c4b",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "import yaml\n",
    "\n",
    "table_descriptor = {\n",
    "    'name': table_name,\n",
    "    'type': 'delta',\n",
    "    'path': f'tables/{table_name}',\n",
    "    'catalog': unity_catalog_name,\n",
    "}\n",
    "\n",
    "# Write table descriptor to lakeFS\n",
    "with branchMain.object(path=f'_lakefs_tables/{table_name}.yaml').writer() as out:\n",
    "    yaml.safe_dump(table_descriptor, out)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "d142bcbd-e16f-4783-9459-c200867e74f0",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "#### Upload the Unity Catalog exporter script to lakeFS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "fac5d0cf-b882-44f5-a531-4f654f73dcea",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "lua_script = \"\"\"\n",
    "\n",
    "local azure = require(\"azure\")\n",
    "local formats = require(\"formats\")\n",
    "local databricks = require(\"databricks\")\n",
    "local delta_exporter = require(\"lakefs/catalogexport/delta_exporter\")\n",
    "local unity_exporter = require(\"lakefs/catalogexport/unity_exporter\")\n",
    "local json = require(\"encoding/json\")\n",
    "\n",
    "local table_descriptors_path = \"_lakefs_tables\"\n",
    "local sc = azure.blob_client(args.azure.storage_account, args.azure.access_key)\n",
    "local function write_object(_, key, buf)\n",
    "  return sc.put_object(key,buf)\n",
    "end\n",
    "local delta_client = formats.delta_client(args.lakefs.access_key_id, args.lakefs.secret_access_key)\n",
    "local delta_table_details = delta_exporter.export_delta_log(action, args.table_defs, write_object, delta_client, table_descriptors_path, azure.abfss_transform_path)\n",
    "\n",
    "-- Register the exported table in Unity Catalog:\n",
    "local databricks_client = databricks.client(args.databricks_host, args.databricks_token)\n",
    "local registration_statuses = unity_exporter.register_tables(action, \"_lakefs_tables\", delta_table_details, databricks_client, args.warehouse_id)\n",
    "for t, status in pairs(registration_statuses) do\n",
    "    print(\"Unity catalog registration for table \\\\\"\" .. t .. \"\\\\\" completed with commit schema status : \" .. status .. \"\\\\n\")\n",
    "end\n",
    "\n",
    "\"\"\"\n",
    "\n",
    "branchMain.object(path=luaScriptName).upload(data=lua_script, mode='wb')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "a93834b0-b4e6-46ca-a912-02fd7cd36857",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "#### Define an action configuration that will run the above script after a commit or merge is completed over the main branch and upload it to lakeFS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "5ad25bb2-1d41-4b6a-ac59-5b21b0b51f85",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "hook_definition = {\n",
    "    'name': 'unity_exporter',\n",
    "    'on': {\n",
    "        'post-commit': {\n",
    "            'branches': [sourceBranch, newBranch+'*']\n",
    "        },\n",
    "        'post-create-branch': {\n",
    "            'branches': [newBranch+'*']\n",
    "        }\n",
    "    },\n",
    "    'hooks': [\n",
    "        {\n",
    "            'id': 'Unity-Registration',\n",
    "            'type': 'lua',\n",
    "            'properties': {\n",
    "                'script_path': luaScriptName,\n",
    "                'args': {\n",
    "                    'azure': {\n",
    "                        'storage_account': '<Azure Storage Account Name>',\n",
    "                        'access_key': 'Azure Storage Account Access Key>',\n",
    "                        'region': '<Azure Region Name>'\n",
    "                    },\n",
    "                    'lakefs': {\n",
    "                        'access_key_id': lakefsAccessKey,\n",
    "                        'secret_access_key': lakefsSecretKey\n",
    "                    },\n",
    "                    'table_defs': [table_name],\n",
    "                    'databricks_host': databricks_host,\n",
    "                    'databricks_token': databricks_token,\n",
    "                    'warehouse_id': warehouse_id\n",
    "                }\n",
    "            }\n",
    "        }\n",
    "    ]\n",
    "}\n",
    "\n",
    "with branchMain.object(path='_lakefs_actions/unity_exporter_action.yaml').writer() as out:\n",
    "    yaml.safe_dump(hook_definition, out)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "91827cc7-c318-4e3a-8692-0390cf626138",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "#### Create the Delta Table in source branch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "899c415e-80b8-468a-a1c3-27d342f6c7dd",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "print(repositoryName)\n",
    "data = [\n",
    "   ('James','Bond','England','intelligence'),\n",
    "   ('Robbie','Williams','England','music'),\n",
    "   ('Hulk','Hogan','USA','entertainment'),\n",
    "   ('Mister','T','USA','entertainment'),\n",
    "   ('Rafael','Nadal','Spain','professional athlete'),\n",
    "   ('Paul','Haver','Belgium','music'),\n",
    "]\n",
    "columns = [\"firstname\",\"lastname\",\"country\",\"category\"]\n",
    "df = spark.createDataFrame(data=data, schema = columns)\n",
    "df.write.option(\"mergeSchema\", \"true\").format(\"delta\").mode(\"overwrite\").partitionBy(\"category\", \"country\").save(f\"lakefs://{repositoryName}/{sourceBranch}/tables/{table_name}\")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "0fb1724f-b393-4cca-995f-871ff0b7cf55",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "#### Commit changes and attach some metadata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "6577b798-3d39-41c3-9a3c-291aa4c0d980",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "branchMain.commit(message='Added configuration files and Delta table!', \n",
    "        metadata={'using': 'python_api'})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "bde9c478-0cc4-4d6d-8b90-3384bd8f9b0a",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "# Demo Starts"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "b930489e-8062-4939-83d0-dab86e75db6c",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "#### Run demo [Setup](./?o=8376305627582670#notebook/866459674872932) tasks here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "8a022372-079c-4d19-a86a-9104c60559da",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "%run ./unityCatalogIntegrationDemoSetup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "cf6c5e7a-c9f1-4f82-bf33-10696e73da32",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "#### Go to SQL Editor, click on \"All\" tab and refresh the schema\n",
    "##### You will notice \"main\" schema under \"lakefs_unity_catalog_demo\" catalog.\n",
    "\n",
    "#### Run SQL to read the data from the main branch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "190c643c-d427-4a14-a609-466960ba88c1",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "df = spark.sql(f\"SELECT * FROM `{unity_catalog_name}`.`{sourceBranch}`.`{table_name}`\")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "88f39fc5-acfe-4a06-8fc3-25be35c8ed10",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "#### Create a new branch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "c469cf81-6f44-4737-8be6-6f4ffdc185aa",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "newBranch = \"dev1\"\n",
    "branchDev = repo.branch(newBranch).create(source_reference=sourceBranch)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "0f74c134-1473-4dc8-9398-7b5c4cbf7b22",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "#### Go back to SQL Editor and refresh the schema\n",
    "##### You will notice new schema for the new branch created in previous step\n",
    "\n",
    "#### Run SQL to read the data from the new branch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "bb470b7d-0b76-4e56-a08f-deaa56777fbd",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "df = spark.sql(f\"SELECT * FROM `{unity_catalog_name}`.`{newBranch}`.`{table_name}`\")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "860d50d7-9d33-47aa-96c2-6cc5bfa0b02f",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "#### Update Delta Table in the new branch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "c4e2e4d6-1b37-483e-a051-cb4d25a0a34d",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col\n",
    "\n",
    "df_us = df.filter(col(\"country\") == \"USA\")\n",
    "df_us.write.format(\"delta\").mode(\"overwrite\").save(f\"lakefs://{repositoryName}/{newBranch}/tables/{table_name}\")\n",
    "df_us.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "4fffd8f9-d113-4a81-9897-4ac778770da3",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "branchDev.commit(message='Updated delta table!', \n",
    "        metadata={'using': 'python_api'})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "1e8b35fc-efee-49b3-8659-8b46b01884e6",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "#### Run SQL to read the data from the new branch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "841971c8-5074-4416-8f94-3a71b4b890fc",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "df = spark.sql(f\"SELECT * FROM `{unity_catalog_name}`.`{newBranch}`.`{table_name}`\")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "80fea7e2-f794-45af-9bcc-bcc0155b23e6",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "#### Run SQL to read the data from the main branch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "398412d9-81ca-45b8-8203-19aa02a1af13",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "df = spark.sql(f\"SELECT * FROM `{unity_catalog_name}`.`{sourceBranch}`.`{table_name}`\")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "90875b7c-41d3-4656-b174-ce74072b455e",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "# Demo Completes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "108a8ed0-234c-4782-8967-fad60dc06d55",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "## More Questions?\n",
    "\n",
    "###### Join the [lakeFS Slack group](https://lakefs.io/slack)"
   ]
  }
 ],
 "metadata": {
  "application/vnd.databricks.v1+notebook": {
   "dashboards": [],
   "language": "python",
   "notebookMetadata": {
    "mostRecentlyExecutedCommandWithImplicitDF": {
     "commandId": 1786061902496977,
     "dataframes": [
      "_sqldf"
     ]
    },
    "pythonIndentUnit": 4
   },
   "notebookName": "Unity Catalog Integration Demo",
   "widgets": {}
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}