{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "cbe0f4e5-e004-4adb-a9b7-d420f438c627",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "# Two tables that must agree\n",
    "\n",
    "You have a ledger and a balance. The balance is just the sum of the postings in the ledger. When\n",
    "both writes go through, they agree. But what if the job dies in between? The ledger moved, the\n",
    "balance did not, and now your books are wrong. No alert, no rollback, nobody notices.\n",
    "\n",
    "This lab walks through that problem, shows how Delta Lake time travel gets you back to a good\n",
    "state, and then shows **catalog-managed tables** with `BEGIN ATOMIC`, which makes the two writes\n",
    "a single unit of work so the gap never happens in the first place.\n",
    "\n",
    "Full instructions: https://brickster.io/labs/two-tables-that-must-agree\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "37363279-dfa8-49ae-bd2c-d7d1a9888bd8",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "## 1. Your salt\n",
    "\n",
    "Copy the eight characters under **Your salt** on the lab page and paste them below. One of the five\n",
    "answers depends on it.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "7df0bd9b-0652-4394-bb61-2122e129c2c7",
     "showTitle": false,
     "startTime": 1790008018369,
     "submitTime": 1790008018369,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "SALT = \"00000000\"  # <- paste yours here\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "63e0a873-cbe1-4e46-a3a9-a1ac43237120",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "## 2. The two tables\n",
    "\n",
    "Unity Catalog managed Delta tables in a schema of your own, built from `spark.range`, so there is\n",
    "nothing to download and the numbers are the same in every workspace.\n",
    "\n",
    "### Why catalog-managed?\n",
    "\n",
    "Normal Delta tables commit on their own. Each insert or merge is its own transaction with its own\n",
    "version. There is no way to say these two writes have to succeed together or not at all. If the job\n",
    "crashes after the first one, you are stuck with half the work done.\n",
    "\n",
    "**Catalog-managed tables** (`delta.feature.catalogManaged = 'supported'`) fix that. They register\n",
    "their commits with the Unity Catalog metastore, which is what makes `BEGIN ATOMIC` work. A\n",
    "multi-statement transaction either commits every change or rolls them all back. And every table in\n",
    "the block has to have the property, so you cannot accidentally mix managed and unmanaged tables.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "finishTime": 1790008055885,
     "inputWidgets": {},
     "nuid": "f5db2586-079a-4962-850e-eef857c52015",
     "showTitle": false,
     "startTime": 1790008021162,
     "submitTime": 1790008021050,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "CATALOG = \"workspace\"  # change this if workspace is not where you can create a schema\n",
    "SCHEMA = \"brickster_two_tables\"\n",
    "T = f\"{CATALOG}.{SCHEMA}\"\n",
    "\n",
    "spark.sql(f\"create schema if not exists {T}\")\n",
    "spark.sql(f\"drop table if exists {T}.ledger\")\n",
    "spark.sql(f\"drop table if exists {T}.balance\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "finishTime": 1790008170688,
     "inputWidgets": {},
     "nuid": "45fbf718-9533-4618-b005-f5a341fdd387",
     "showTitle": false,
     "startTime": 1790008163664,
     "submitTime": 1790008163600,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "spark.sql(f\"\"\"\n",
    "  create table {T}.ledger (id bigint, account int, amount int, batch int)\n",
    "  using delta\n",
    "  tblproperties ('delta.feature.catalogManaged' = 'supported')\n",
    "\"\"\")\n",
    "\n",
    "spark.sql(f\"\"\"\n",
    "  create table {T}.balance (account int, amount bigint)\n",
    "  using delta\n",
    "  tblproperties ('delta.feature.catalogManaged' = 'supported')\n",
    "\"\"\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "0ccad11e-7920-4fa5-9522-bb862b986519",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "The opening ledger is 4,096 postings over 64 accounts. The balance is derived from it, so the two\n",
    "agree the moment they are written. This is the starting point. The rest of the lab will break this\n",
    "agreement and then fix it.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "finishTime": 1790008189717,
     "inputWidgets": {},
     "nuid": "53339bf7-a0b5-4ae8-9e47-729cc30f6032",
     "showTitle": false,
     "startTime": 1790008175169,
     "submitTime": 1790008175135,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "spark.sql(f\"\"\"\n",
    "  insert into {T}.ledger\n",
    "  select id, cast(id % 64 as int), cast((id * 37) % 101 - 50 as int), 0\n",
    "  from range(0, 4096)\n",
    "\"\"\")\n",
    "\n",
    "spark.sql(f\"\"\"\n",
    "  insert into {T}.balance\n",
    "  select account, sum(amount) from {T}.ledger group by account\n",
    "\"\"\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "bda4f6c1-dca1-47f0-afc0-530f171017c6",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "### The invariant\n",
    "\n",
    "The one question this lab keeps asking: does each account's balance match the sum of its postings?\n",
    "Right now every account does. The rest of the lab is about what it takes to keep that true when a\n",
    "job fails halfway through.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "finishTime": 1790008232475,
     "inputWidgets": {},
     "nuid": "c2e24c6d-bdde-4fd6-b81d-d73f1cccf0aa",
     "showTitle": false,
     "startTime": 1790008229860,
     "submitTime": 1790008229812,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "def disagreeing():\n",
    "    return spark.sql(f\"\"\"\n",
    "      select count(*) as n\n",
    "      from {T}.balance b\n",
    "      join (select account, sum(amount) as s from {T}.ledger group by account) l\n",
    "        on l.account = b.account\n",
    "      where b.amount <> l.s\n",
    "    \"\"\").collect()[0][\"n\"]\n",
    "\n",
    "\n",
    "print(\"accounts disagreeing:\", disagreeing())\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "59cee794-e93a-422c-a301-0dc5e49f20df",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "## 3. The late batch, the way the job runs it today\n",
    "\n",
    "Here is the failure. A batch job posts 300 new rows to the ledger, then in a separate statement\n",
    "updates the balance. On Tuesday the job died between the two. Each statement commits on its own,\n",
    "so there is nothing to roll back. The ledger has the new rows, the balance does not. Every report\n",
    "that reads balances is now wrong, and no alert fired because both tables look fine on their own.\n",
    "\n",
    "The cell below replays that. It runs the first statement and stops, just like the crash. Nothing\n",
    "rolls back because nothing was ever a transaction.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "finishTime": 1790008261398,
     "inputWidgets": {},
     "nuid": "d1d21746-d049-4bd5-9f5c-1fb8c7b68009",
     "showTitle": false,
     "startTime": 1790008258422,
     "submitTime": 1790008258385,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "spark.sql(f\"\"\"\n",
    "  insert into {T}.ledger\n",
    "  select 10000 + id, cast((id * 13 + 5) % 47 as int), cast((id * 29) % 61 - 30 as int), 1\n",
    "  from range(0, 300)\n",
    "\"\"\")\n",
    "\n",
    "# The update to balance never ran.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "74f04e7b-8399-4497-8042-5136a8f9bf34",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "### How wide is the damage\n",
    "\n",
    "The failed batch touched accounts spread across the table, so the disagreement is not just one\n",
    "account. This count tells us how many balances are now wrong. It is the kind of number a\n",
    "reconciliation job would catch hours later, long after the person who could have fixed it has gone\n",
    "home.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: count the accounts whose balance no longer matches their ledger.\n",
    "drift_accounts = None\n",
    "\n",
    "print(f\"BRICKSTER:drift_accounts={drift_accounts}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "6d171896-ff8b-43bc-bac2-30356d0d6f0b",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "## 4. Back to a state you can trust\n",
    "\n",
    "The ledger moved and the balance did not. Before we fix the process, we have to undo the damage.\n",
    "\n",
    "**Time travel** is Delta Lake's answer. Every write creates a new table version, and old versions stick\n",
    "around. So you can `RESTORE TABLE ... TO VERSION AS OF` the version just before the failed insert.\n",
    "No need to figure out which 300 rows were yours and delete them by hand. `DESCRIBE HISTORY` lists\n",
    "every version so you know which one was the last good one.\n",
    "\n",
    "This is the safety net that plain file-based systems do not have. No backup, no restore script, no\n",
    "row-level diff. One command takes you back.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: put the ledger back the way it was before the failed batch.\n",
    "# `describe history` lists every version, and `restore table ... to version as of <n>` goes back\n",
    "# to one. Nothing here needs to know which 300 rows were yours.\n",
    "\n",
    "print(\"accounts disagreeing:\", disagreeing())\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "12103245-26d1-4bb0-988e-217e367f252e",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "## 5. Replay it atomically\n",
    "\n",
    "Now we replay the same batch, but this time the ledger insert and the balance merge are inside a\n",
    "single `BEGIN ATOMIC ... END` block. Both tables are catalog-managed, so the metastore treats them\n",
    "as one transaction. Either both writes commit together, or neither does.\n",
    "\n",
    "### Why this matters\n",
    "\n",
    "This is the fix for the failure we just saw. If the job crashes after the insert but before the\n",
    "merge, or if the merge itself fails, the atomic block rolls back the insert for you. The ledger\n",
    "and the balance are never left out of sync. No time-travel rescue, no manual cleanup, no 3 AM phone\n",
    "call.\n",
    "\n",
    "`BEGIN ATOMIC ... END;` is the form to use here: serverless compute runs non-interactive\n",
    "transactions only, so `BEGIN TRANSACTION; ... COMMIT;` belongs to a SQL warehouse and not to this\n",
    "notebook. If your runtime will not take the compound statement through `spark.sql`, run the same\n",
    "block in a `%sql` cell and carry on.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: post the same 300 rows and bring the balance with them, both inside one\n",
    "# `begin atomic ... end` block.\n",
    "#\n",
    "# One quirk worth knowing before you start: inside a transaction block the `range()` table\n",
    "# function is not available. Generate the rows with `explode(sequence(0, 299)) as t(id)` instead.\n",
    "\n",
    "print(\"accounts disagreeing:\", disagreeing())\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "e7190f65-e4e9-429e-89b4-6a82619410bc",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "## 6. What a transaction that fails leaves behind\n",
    "\n",
    "Time travel got us out of trouble in Section 4, but we had to *know* the table was wrong and *find*\n",
    "the right version. Atomic transactions make that unnecessary. They fail cleanly.\n",
    "\n",
    "Here we deliberately abort a transaction. We insert ten rows into the ledger, then `SIGNAL` an error\n",
    "inside the same `BEGIN ATOMIC` block. The insert runs, the signal fires, and the whole block rolls\n",
    "back. The questions are: did the ledger gain a version? Did it gain rows?\n",
    "\n",
    "The error is not the answer. What the ledger holds afterwards is. And so is what `describe history`\n",
    "recorded or did not record while the block was running. If the answer is zero new versions and zero\n",
    "new rows, that is the whole point. A failed atomic transaction leaves nothing behind.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "finishTime": 1790009073494,
     "inputWidgets": {},
     "nuid": "8b0eb866-08c1-41bb-9a86-a6d7e1921dac",
     "showTitle": false,
     "startTime": 1790009069471,
     "submitTime": 1790009069436,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "versions_before = spark.sql(f\"describe history {T}.ledger\").count()\n",
    "\n",
    "try:\n",
    "    spark.sql(f\"\"\"\n",
    "      begin atomic\n",
    "        insert into {T}.ledger\n",
    "        select 20000 + id, cast(id % 64 as int), 1, 2 from explode(sequence(0, 9)) as t(id);\n",
    "\n",
    "        if (select count(*) from {T}.ledger where batch = 2) > 0 then\n",
    "          signal sqlstate '45000' set message_text = 'batch 2 was never meant to land';\n",
    "        end if;\n",
    "      end\n",
    "    \"\"\")\n",
    "    print(\"no error, which is not what this cell is for\")\n",
    "except Exception as e:\n",
    "    print(\"refused:\", type(e).__name__)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "finishTime": 1790009146794,
     "inputWidgets": {},
     "nuid": "b493dfa8-8279-480e-9fd0-d12f0be23f15",
     "showTitle": false,
     "startTime": 1790009141018,
     "submitTime": 1790009140964,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "display(spark.sql(f\"describe history {T}.ledger\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: how many versions did the ledger gain while that ran, and how many rows does it hold now.\n",
    "abort_versions = None\n",
    "abort_rows = None\n",
    "\n",
    "print(f\"BRICKSTER:abort_versions={abort_versions}\")\n",
    "print(f\"BRICKSTER:abort_rows={abort_rows}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "a9dc8668-bc86-44d9-aa14-b085726a9a65",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "## 7. The last two answers\n",
    "\n",
    "The final check ties it all together. `mismatch_after` confirms that the atomic replay in Section 5\n",
    "left the tables in agreement. Zero disagreeing accounts, same as before the crash. The per-account\n",
    "probes verify the balances themselves are correct.\n",
    "\n",
    "If every step above landed, these numbers are right. If any step failed silently, they would not\n",
    "be. That is the guarantee catalog-managed atomic transactions give you: either the whole batch is in\n",
    "the books, or none of it is.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mismatch_after = disagreeing()\n",
    "print(f\"BRICKSTER:mismatch_after={mismatch_after}\")\n",
    "\n",
    "# Your salt picks the account. Somebody else's notebook answers this with their account.\n",
    "bucket = int(SALT, 16) % 64\n",
    "probe = spark.sql(f\"select amount from {T}.balance where account = {bucket}\").collect()[0][\"amount\"]\n",
    "print(f\"BRICKSTER:probe={probe}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Check your run\n",
    "\n",
    "Two ways, same grader.\n",
    "\n",
    "**Export and upload.** *File > Export > IPYNB*, then drop it on the lab page.\n",
    "\n",
    "**Or straight from here.** Make a token on https://brickster.io/settings, put it in the `brickster`\n",
    "secret scope once per workspace, and run the last cell. Your Databricks credentials never leave\n",
    "this workspace.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install brickster\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import brickster\n",
    "\n",
    "brickster.submit(\"two-tables-that-must-agree\")\n"
   ]
  }
 ],
 "metadata": {
  "application/vnd.databricks.v1+notebook": {
   "computePreferences": null,
   "dashboards": [],
   "environmentMetadata": {
    "base_environment": "",
    "environment_version": "5"
   },
   "inputWidgetPreferences": null,
   "language": "python",
   "notebookMetadata": {
    "pythonIndentUnit": 4
   },
   "notebookName": "starter",
   "widgets": {}
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
