{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "6e11b9b6-c8ac-4e68-a21d-e68154e79c0d",
   "metadata": {},
   "source": [
    "# Import from GCS\n",
    "\n",
    "Elevate Data empowers organizations to effectively manage their data lifecycle by enabling seamless integration of diverse data sources on a large scale, while also providing tools for enhanced governance and efficiency.\n",
    "\n",
    "Elevate Data manages connections to various data sources (S3, Snowflake, BigQuery and GCS at this stage), allowing for efficient data extraction.\n",
    "\n",
    "It also provides a central repository of data elements, imported from different sources, for greater transparency and efficient collaboration.\n",
    "\n",
    "This notebook outlines how to: \n",
    "\n",
    "1. Create a Connection: Connections allow for a secure setup to external data sources, such as a Snowflake host, a Databricks Delta Sharing, a GCS bucket or an Amazon S3 bucket.\n",
    "2. Retrieve and list connections to manage your data sources.\n",
    "3. Create a Data Source: A Data Source defines the specific elements within the connection from which data is extracted.\n",
    "4. Create a Data Table: A Data Table references a Data Source and triggers the import of its first version. With `blocking=True`, the call waits for that import to finish before returning.\n",
    "5. Retrieve and list data tables to manage your data assets.\n",
    "6. Preview 100 rows of the version.\n",
    "7. Create additional data versions to refresh your data.\n",
    "\n",
    "In this example, we are going to create a Connection to a GCS and import data from it using a Data Table and create a new version of the data table."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a344b11a-7b49-45de-a86c-52f6fd1a2598",
   "metadata": {},
   "source": [
    "## Install the SDK"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d57ae996-2b78-4cc9-ac57-84d48302a4b8",
   "metadata": {
    "tags": [
     "skip-execution"
    ]
   },
   "outputs": [],
   "source": [
    "%pip install --upgrade earnix-elevate"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6b60e2d-354c-430e-809e-45cb1dc89126",
   "metadata": {},
   "source": [
    "## Set your Elevate server and credentials\n",
    "\n",
    "In this example we use environment variables to authenticate to Elevate, but you can also inject your credentials using Python arguments to each Service's client, like in the commented example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72923393-52af-43f9-b7fe-8a4d24fff0d9",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-04-09T15:11:37.745173Z",
     "iopub.status.busy": "2026-04-09T15:11:37.744901Z",
     "iopub.status.idle": "2026-04-09T15:11:37.754885Z",
     "shell.execute_reply": "2026-04-09T15:11:37.754475Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ.setdefault(\"E2_SERVER\", \"YOUR_SERVER\")\n",
    "os.environ.setdefault(\"E2_CLIENT_ID\", \"YOUR_CLIENT_ID\")\n",
    "os.environ.setdefault(\"E2_SECRET_KEY\", \"YOUR_SECRET_KEY\")\n",
    "\n",
    "# Alternatively:\n",
    "# ConnectionService(\n",
    "#     server=\"YOUR_SERVER\",\n",
    "#     client_id=\"YOUR_CLIENT_ID\",\n",
    "#     secret_key=\"YOUR_SECRET_KEY\"\n",
    "# )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83e2ea7f-29d5-49f4-a234-8ca67562a6cd",
   "metadata": {},
   "source": [
    "## Imports and demo preparation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7cf4160-db73-4924-8e2d-a61cb511208a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-04-09T15:11:37.756739Z",
     "iopub.status.busy": "2026-04-09T15:11:37.756563Z",
     "iopub.status.idle": "2026-04-09T15:11:39.661243Z",
     "shell.execute_reply": "2026-04-09T15:11:39.660579Z"
    }
   },
   "outputs": [],
   "source": [
    "import json\n",
    "from datetime import datetime\n",
    "\n",
    "from earnix_elevate import (\n",
    "    ConnectionService,\n",
    "    CreateBucketDataSourceRequest,\n",
    "    CreateDataTableRequest,\n",
    "    CreateGcsConnectionRequest,\n",
    "    DataSourceService,\n",
    "    DataTableService,\n",
    "    GcpWorkloadIdentityAuthRequest,\n",
    ")\n",
    "\n",
    "DEMO_SUFFIX = datetime.now().strftime(\"%Y-%m-%d %H:%M:%S\")\n",
    "\n",
    "connection_service = ConnectionService()\n",
    "data_source_service = DataSourceService()\n",
    "data_table_service = DataTableService()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "32d3a652-fbf0-4167-a571-43c535f6127b",
   "metadata": {},
   "source": [
    "## Create a Connection\n",
    "\n",
    "There are several steps to follow in order to create a GCS connection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f33d429-6381-4341-9642-b9b9352282ab",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-04-09T15:11:39.664140Z",
     "iopub.status.busy": "2026-04-09T15:11:39.663922Z",
     "iopub.status.idle": "2026-04-09T15:11:40.608988Z",
     "shell.execute_reply": "2026-04-09T15:11:40.608196Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "service_account_file_content = json.dumps(\n",
    "    {\n",
    "        # file content example:\n",
    "        #   \"type\": \"service_account\",\n",
    "        #   \"project_id\": \"projectId\",\n",
    "        #   \"private_key_id\": \"b571111bd64a3081a11111870827664c111111\",\n",
    "        #   \"private_key\": (\n",
    "        #       \"-----BEGIN PRIVATE KEY--...\n",
    "    }\n",
    ")\n",
    "\n",
    "# Option 1: OAuth2 Client Credentials\n",
    "auth_details = GcpWorkloadIdentityAuthRequest(\n",
    "    authType=\"gcp_wif\",\n",
    "    clientId=os.environ.get(\"GCP_CLIENT_ID\", \"YOUR_CLIENT_ID\"),\n",
    "    clientSecret=os.environ.get(\"GCP_CLIENT_SECRET\", \"YOUR_CLIENT_SECRET\"),\n",
    "    tokenUrl=os.environ.get(\"GCP_TOKEN_URL\", \"YOUR_GCP_TOKEN_URL\"),\n",
    "    audience=os.environ.get(\"GCP_AUDIENCE\", \"YOUR_GCP_AUDIENCE\"),\n",
    "    scope=os.environ.get(\"GCP_SCOPE\", \"YOUR_GCP_SCOPE\"),\n",
    "    serviceAccountEmail=os.environ.get(\n",
    "        \"GCP_SERVICE_ACCOUNT_EMAIL\", \"YOUR_GCP_SERVICE_ACCOUNT_EMAIL\"\n",
    "    ),\n",
    "    projectId=os.environ.get(\"GCP_PROJECT_ID\", \"YOUR_GCP_PROJECT_ID\"),\n",
    ")\n",
    "\n",
    "# Option 2: Service Account Key\n",
    "# auth_details = GcpServiceAccountKeyAuthRequest(\n",
    "#    authType=\"gcp_sak\",\n",
    "#    projectId=\"xxxxx\",\n",
    "#    serviceAccountKeyJson=service_account_file_content,\n",
    "# )\n",
    "\n",
    "new_conn = CreateGcsConnectionRequest(\n",
    "    type=\"gcs\",\n",
    "    name=\"GCS Demo Conn \" + DEMO_SUFFIX,\n",
    "    description=\"GCS connection example\",\n",
    "    auth=auth_details,\n",
    ")\n",
    "\n",
    "create_conn = connection_service.create_connection(new_conn)\n",
    "print(f\"\\n{create_conn=}\")\n",
    "\n",
    "# connection_service.test_connection_by_id(create_conn.id);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ff0487b-3abc-420e-8812-1ce00a375b23",
   "metadata": {},
   "source": [
    "### Retrieve Connection Details\n",
    "\n",
    "After creating a connection, you can retrieve its details using the connection ID."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f50fa357-14f9-4854-8773-6784d2522d47",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-04-09T15:11:40.611454Z",
     "iopub.status.busy": "2026-04-09T15:11:40.611254Z",
     "iopub.status.idle": "2026-04-09T15:11:40.894236Z",
     "shell.execute_reply": "2026-04-09T15:11:40.893786Z"
    }
   },
   "outputs": [],
   "source": [
    "retrieved_conn = connection_service.get_connection(create_conn.id)\n",
    "print(f\"\\n{retrieved_conn=}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ea9bb93-3ca6-4b22-9f80-4acac695e295",
   "metadata": {},
   "source": [
    "### List All Connections\n",
    "\n",
    "You can also list all available connections in your environment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20432ffd-6085-4acc-9bd7-a02cc2f10566",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-04-09T15:11:40.896258Z",
     "iopub.status.busy": "2026-04-09T15:11:40.896036Z",
     "iopub.status.idle": "2026-04-09T15:11:42.689756Z",
     "shell.execute_reply": "2026-04-09T15:11:42.689389Z"
    }
   },
   "outputs": [],
   "source": [
    "all_connections = connection_service.list_connections().items\n",
    "print(f\"\\nFound {len(all_connections)} connections:\")\n",
    "for conn in all_connections:\n",
    "    print(f\"  - {conn.name} (ID: {conn.id})\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "9172391d-4a3d-445e-902a-f670ac67388c",
   "metadata": {},
   "source": [
    "## Create a Data Table\n",
    "\n",
    "There are several steps to follow in order to create a Data Table that references specific GCS elements through the established connection. Once the Data Table is created, the import of the first version will begin automatically."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "295a8faa-c61d-44f8-bad9-549791f725c2",
   "metadata": {},
   "source": [
    "### Create the Data Source \n",
    "\n",
    "When creating a Data Source for the Data Table, provide the ID of the connection you wish to use, and define the table name, schema, and share to retrieve data from the source."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ef529f2-1eee-4f27-9344-12d3e2052c70",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-04-09T15:11:42.691062Z",
     "iopub.status.busy": "2026-04-09T15:11:42.690978Z",
     "iopub.status.idle": "2026-04-09T15:11:43.306538Z",
     "shell.execute_reply": "2026-04-09T15:11:43.305690Z"
    }
   },
   "outputs": [],
   "source": [
    "new_ds = CreateBucketDataSourceRequest(\n",
    "    type=\"bucket\",\n",
    "    name=\"GCS Demo DS \" + DEMO_SUFFIX,\n",
    "    connectionId=create_conn.id,\n",
    "    path=\"gs://elevate_dev_bucket/ALDF_PRIVATE/boats/boats.csv\",\n",
    "    description=\"Test GCS DataSource for data source tests\",\n",
    ")\n",
    "\n",
    "create_ds = data_source_service.create_data_source(new_ds)\n",
    "print(f\"\\n{create_ds=}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "02ec71cf-471c-4dee-9a7d-92447c9eebad",
   "metadata": {},
   "source": [
    "### Create the Data Table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06492544-095d-4dbc-a2ae-843e06dcc8ae",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-04-09T15:11:43.311546Z",
     "iopub.status.busy": "2026-04-09T15:11:43.311180Z",
     "iopub.status.idle": "2026-04-09T15:11:44.116719Z",
     "shell.execute_reply": "2026-04-09T15:11:44.115416Z"
    }
   },
   "outputs": [],
   "source": [
    "new_dt = CreateDataTableRequest(\n",
    "    name=\"Bee-Insurance Demo DT \" + DEMO_SUFFIX,\n",
    "    dataSourceId=create_ds.id,\n",
    ")\n",
    "\n",
    "# blocking=True creates the table and waits for its import to finish,\n",
    "# returning the imported table (raising on failure/timeout).\n",
    "create_dt = data_table_service.create_data_table(new_dt, blocking=True)\n",
    "print(f\"\\n{create_dt=}\")\n",
    "data_version_number = create_dt.latest_data_version.version_number"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "df46f2c8-3bd2-43b7-82b3-01018454b670",
   "metadata": {},
   "source": [
    "### Retrieve Data Table Details\n",
    "\n",
    "After creating a data table, you can retrieve its details using the data table ID."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06f37cae-396c-4a59-8ef6-801502956000",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-04-09T15:11:44.121219Z",
     "iopub.status.busy": "2026-04-09T15:11:44.120776Z",
     "iopub.status.idle": "2026-04-09T15:11:44.307370Z",
     "shell.execute_reply": "2026-04-09T15:11:44.306314Z"
    }
   },
   "outputs": [],
   "source": [
    "retrieved_dt = data_table_service.get_data_table(create_dt.id)\n",
    "print(f\"\\n{retrieved_dt=}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e5191dfc-62f8-46cb-8077-540eb3808904",
   "metadata": {},
   "source": [
    "### List All Data Tables\n",
    "\n",
    "You can also list all available data tables in your environment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e6ebbee-bd3c-44f7-b88b-54defc0f6fdb",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-04-09T15:11:44.313344Z",
     "iopub.status.busy": "2026-04-09T15:11:44.312968Z",
     "iopub.status.idle": "2026-04-09T15:11:47.731351Z",
     "shell.execute_reply": "2026-04-09T15:11:47.731090Z"
    }
   },
   "outputs": [],
   "source": [
    "all_data_tables = data_table_service.list_data_tables().items\n",
    "print(f\"\\nFound {len(all_data_tables)} data tables:\")\n",
    "for dt in all_data_tables:\n",
    "    print(f\"  - {dt.name} (ID: {dt.id})\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1a93867-4041-4aa4-9856-c874e6dbc99b",
   "metadata": {},
   "source": [
    "## Preview\n",
    "\n",
    "Once the version has been successfully imported, you can preview the first 100 rows of the imported data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e855bb0c-6f29-4504-bfa4-fb7e3d6ec617",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-04-09T15:14:08.692257Z",
     "iopub.status.busy": "2026-04-09T15:14:08.691885Z",
     "iopub.status.idle": "2026-04-09T15:14:09.271769Z",
     "shell.execute_reply": "2026-04-09T15:14:09.270396Z"
    }
   },
   "outputs": [],
   "source": [
    "dvt = data_table_service.get_data_version_preview(create_dt.id, data_version_number)\n",
    "print(f\"\\n{dvt=}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f995022a-044b-4f59-b526-eb03f603c1de",
   "metadata": {},
   "source": [
    "## Create Additional Data Versions\n",
    "\n",
    "You can create new versions of your data table to retrieve updated data from the source. This is useful when the underlying data in Databricks has changed and you want to refresh your data table."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8f278a9-5788-4f6a-ad04-f63aab863840",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-04-09T15:14:09.276130Z",
     "iopub.status.busy": "2026-04-09T15:14:09.275771Z",
     "iopub.status.idle": "2026-04-09T15:14:09.629812Z",
     "shell.execute_reply": "2026-04-09T15:14:09.628815Z"
    }
   },
   "outputs": [],
   "source": [
    "new_version = data_table_service.create_data_version(create_dt.id)\n",
    "print(f\"\\n{new_version=}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cleanup-databricks-header",
   "metadata": {},
   "source": [
    "## Cleanup\n",
    "\n",
    "Delete all resources created in this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cleanup-databricks",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-04-09T15:14:09.633607Z",
     "iopub.status.busy": "2026-04-09T15:14:09.633308Z",
     "iopub.status.idle": "2026-04-09T15:14:12.989472Z",
     "shell.execute_reply": "2026-04-09T15:14:12.987954Z"
    }
   },
   "outputs": [],
   "source": [
    "# Cleanup — delete all resources created in this notebook\n",
    "try:\n",
    "    data_table_service.delete_data_table(data_table_id=create_dt.id)\n",
    "    print(f\"Deleted data table {create_dt.id}\")\n",
    "except Exception as e:\n",
    "    print(f\"Failed to delete data table: {e}\")\n",
    "\n",
    "try:\n",
    "    data_source_service.delete_data_source(data_source_id=create_ds.id)\n",
    "    print(f\"Deleted data source {create_ds.id}\")\n",
    "except Exception as e:\n",
    "    print(f\"Failed to delete data source: {e}\")\n",
    "\n",
    "try:\n",
    "    connection_service.delete_connection(connection_id=create_conn.id)\n",
    "    print(f\"Deleted connection {create_conn.id}\")\n",
    "except Exception as e:\n",
    "    print(f\"Failed to delete connection: {e}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}