Add metadata on a Notebook-basis in Databricks?

[diegoeche]

First, our team just tested 0.6.4. We <3 <3 the support for delta! 👯

I'm not an expert in Databricks btw, just somebody that can write a little Scala.

So we have this problem. We have multiple notebooks in Databricks, and we would like to attach the notebook name to the extra properties of Spline.

In Databricks we have something like this:

import os
token = os.environ["NOTEBOOK_PATH"] = dbutils.notebook.entry_point.getDbutils().notebook().getContext().notebookPath().get()
print("Configured: ", os.environ.get("NOTEBOOK_PATH"))

I thought that maybe I could set an env variable and then read it in the filter:

class EnvFilter(configuration: Configuration) extends PostProcessingFilter {
  def createExtraProperties() = {
    val notebookName = sys.env("NOTEBOOK_PATH")
    Map(
      "notebookName" -> notebookName
    )
  }
  override def processExecutionEvent(event: ExecutionEvent, ctx: HarvestingContext): ExecutionEvent =
    event.withAddedExtra(createExtraProperties())
  override def processExecutionPlan(plan: ExecutionPlan, ctx: HarvestingContext ): ExecutionPlan =
    plan.withAddedExtra(createExtraProperties())
  override def processReadOperation(op: ReadOperation, ctx: HarvestingContext ): ReadOperation =
    op.withAddedExtra(createExtraProperties())
  override def processWriteOperation(op: WriteOperation, ctx: HarvestingContext): WriteOperation =
    op.withAddedExtra(createExtraProperties())
  override def processDataOperation(op: DataOperation, ctx: HarvestingContext  ): DataOperation =
    op.withAddedExtra(createExtraProperties())
}

But the value is always empty. I also tried writing it into the config but is also not in the config. I'm using declarative initialization. I guess I could move to programmatic initialization and set the value then.

I guess, my general question. In Databricks, how can you "set" or "pass" information to the filter since the filter only has access to the Configuration in the constructor. Is there like a singleton somewhere else to read these values?

Thanks in advance! And thanks for this great project.

[diegoeche]

Maybe to provide a little more context.

I tried this:

import com.databricks.dbutils_v1.DBUtilsHolder.dbutils
class DBUtilsFilter(configuration: Configuration) extends PostProcessingFilter {
  def createExtraProperties() = {
    Map(
      "notebookName" -> dbutils.notebook.getContext.notebookPath
    )
  }
  override def processExecutionEvent(event: ExecutionEvent, ctx: HarvestingContext): ExecutionEvent =
    event.withAddedExtra(createExtraProperties())
  override def processExecutionPlan(plan: ExecutionPlan, ctx: HarvestingContext ): ExecutionPlan =
    plan.withAddedExtra(createExtraProperties())
  override def processReadOperation(op: ReadOperation, ctx: HarvestingContext ): ReadOperation =
    op.withAddedExtra(createExtraProperties())
  override def processWriteOperation(op: WriteOperation, ctx: HarvestingContext): WriteOperation =
    op.withAddedExtra(createExtraProperties())
  override def processDataOperation(op: DataOperation, ctx: HarvestingContext  ): DataOperation =
    op.withAddedExtra(createExtraProperties())
}

Compiled it with this SBT:

scalaVersion := "2.12.14"

name := "dataopsfilter"
organization := "com.myorg"
version := "1.0"

// You can define other libraries as dependencies in your build like this:
libraryDependencies += "org.scala-lang.modules" %% "scala-parser-combinators" % "1.1.2"
libraryDependencies += "za.co.absa.spline.agent.spark" %% "agent-core" % "0.6.2"
libraryDependencies += "com.databricks" % "dbutils-api_2.12" % "0.0.5"

Error:

21/09/20 20:39:08 ERROR SplineQueryExecutionListener: Unexpected error occurred during lineage processing for application: Databricks Shell #local-1632170264454
java.lang.NullPointerException
	at sun.reflect.NativeMethodAccessorImpl.invoke0(Native Method)
	at sun.reflect.NativeMethodAccessorImpl.invoke(NativeMethodAccessorImpl.java:62)
	at sun.reflect.DelegatingMethodAccessorImpl.invoke(DelegatingMethodAccessorImpl.java:43)
	at java.lang.reflect.Method.invoke(Method.java:498)
	at com.databricks.dbutils_v1.DBUtilsHolder$$anon$1.invoke(DBUtilsHolder.scala:17)
	at com.sun.proxy.$Proxy74.notebook(Unknown Source)
	at za.co.absa.spline.harvester.postprocessing.DBUtilsFilter.createExtraProperties(Main.scala:72)

[wajda]

I think the env variable doesn't work because the JVM process might be already running at the time you set the env var in the notebook, and thus one isn't visible to Spline.

I'm not sure about dbutils though (I don't know how DBUtilsHolder is implemented).

Programmatic initialization would hardly help here because neither of those cases is related to the time of calling a filter constructor.

Try using system properties for the purpose instead of environment variables.

// in the notebook
sys.props("NOTEBOOK_PATH") = ...

// in the `createExtraProperties()` method
val notebookPath = sys.props("NOTEBOOK_PATH")

[cerveada]

One note about the configuration:

The Configuration injected in the filter via constructor is always only subset of the configuration that have the filter namespace.

For example, this is configuration for DataSourcePasswordReplacingFilter:

spline.postProcessingFilter.dsPasswordReplace.className=za.co.absa.spline.harvester.postprocessing.DataSourcePasswordReplacingFilter
spline.postProcessingFilter.dsPasswordReplace.replacement=*****

and in the filter constructor I would get the replacement string like this:

def this(conf: Configuration) = this(
  conf.getRequiredString("replacement")
)

In other words: only keys starting with spline.postProcessingFilter.dsPasswordReplace. are available inside the filter.

You can see the configuration in spline.default.properties and the filter in DataSourcePasswordReplacingFilter.scala

The configuration must be set before the Filter is initialized.

[diegoeche]

Thanks so much! Yes, properties worked just fine!

I did something like this in the filter:

class PropertyFilter(configuration: Configuration) extends PostProcessingFilter {
  def createExtraProperties() = {
    Map(
      "notebookName" -> System.getProperty("spline.config.notebook")
    )
  }
  override def processExecutionEvent(event: ExecutionEvent, ctx: HarvestingContext): ExecutionEvent =
    event.withAddedExtra(createExtraProperties())
  override def processExecutionPlan(plan: ExecutionPlan, ctx: HarvestingContext ): ExecutionPlan =
    plan.withAddedExtra(createExtraProperties())
  override def processReadOperation(op: ReadOperation, ctx: HarvestingContext ): ReadOperation =
    op.withAddedExtra(createExtraProperties())
  override def processWriteOperation(op: WriteOperation, ctx: HarvestingContext): WriteOperation =
    op.withAddedExtra(createExtraProperties())
  override def processDataOperation(op: DataOperation, ctx: HarvestingContext  ): DataOperation =
    op.withAddedExtra(createExtraProperties())
}

And in the notebook I have a command that sets the property like this:

%scala
System.setProperty("spline.config.notebook", dbutils.notebook.getContext.notebookPath.mkString)

[wajda]

Do you use the way described here?
https://github.com/AbsaOSS/spline-getting-started/blob/main/spline-on-databricks/README.md

[DhananjayPunekar]

HI @wajda , I have tried with various run time versions of databricks and it seems to be working fine for me. Currently I have codeless initialisation enabled on databricks runtime 9.1 LTS.

Yes, I use almost similar method described in the documentation link that you have pasted. The major difference is that, I a not installing JAR manually. I am not setting the spark configs manually. I have created a .sh file, which includes the JAR installation and spark config setup. This .sh file can be set as a databricks global init script (to run on all cluster startups in the databricks account) or cluster-scoped init script (to run on specific clusters only)

The .sh file that I use is:

#!/bin/bash

wget --quiet https://repo1.maven.org/maven2/za/co/absa/spline/agent/spark/spark-3.1-spline-agent-bundle_2.12/0.7.1/spark-3.1-spline-agent-bundle_2.12-0.7.1.jar -O /databricks/jars/spark-3.1-spline-agent-bundle_2.12-0.7.1.jar

cat <<EOF > /databricks/driver/conf/00-custom-spark-driver-defaults.conf
[driver] {
"spark.spline.mode" = "REQUIRED",
"spark.spline.producer.url" = azure-func-app-url,
"spark.sql.queryExecutionListeners" = "za.co.absa.spline.harvester.listener.SplineQueryExecutionListener"
}
EOF

The main advantage of having this in the init script (as a .sh) is that, the JAR is installed and spark confs are set along with databricks start up.

In this repo, they have the code for how the Azure Function App has to be set up as the Spline Gateway. https://github.com/intellishore/data-lineage-databricks-to-purview

[DhananjayPunekar]

The class has to be on the driver class path. Then you simply register it via properties, no matter what initialization type you are using.

May be I am just a noob, but I still do not understand this completely. Which class has to be on the driver class path? How to register it via properties? Is there any demo of how this can be done?

[wajda]

You can use the same approach as you use to install the Spline agent to the Spark driver – compile a custom filter class, put it into the JAR and copy to the /databricks/jars folder.
See the example project here - https://github.com/AbsaOSS/spline-getting-started/tree/main/spark-agent-extension-example

[zacayd]

HI @diegoeche
I saw you are member of the spine project
I am struggling to inject the notebook information inside the execution plan they produce
Till know I did it by addition each notebook a scala code,
But I am trying to get rid of it and do it in one centralize place
Do in any chance you succeeded to do it (like codless,AzureFunction)?
Help will be much appreciated
Zacay

%scala
import scala.util.parsing.json.JSON
import za.co.absa.spline.harvester.SparkLineageInitializer._
import za.co.absa.spline.agent.AgentConfig
import za.co.absa.spline.harvester.postprocessing.AbstractPostProcessingFilter
import za.co.absa.spline.harvester.postprocessing.PostProcessingFilter
import org.apache.commons.configuration.Configuration
import za.co.absa.spline.harvester.conf.StandardSplineConfigurationStack
import za.co.absa.spline.harvester.HarvestingContext
import za.co.absa.spline.producer.model.ExecutionPlan
import za.co.absa.spline.producer.model.ExecutionEvent
import za.co.absa.spline.producer.model.ReadOperation
import za.co.absa.spline.producer.model.WriteOperation
import za.co.absa.spline.producer.model.DataOperation
import za.co.absa.spline.harvester.ExtraMetadataImplicits._
import za.co.absa.spline.harvester.SparkLineageInitializer._

val notebookInformationJson = dbutils.notebook.getContext.toJson
val outerMap = JSON.parseFull(notebookInformationJson).getOrElse(0).asInstanceOf[Map[String,String]]
val tagMap = outerMap("tags").asInstanceOf[Map[String,String]]
val extraContextMap = outerMap("extraContext").asInstanceOf[Map[String,String]]
val notebookPath = extraContextMap("notebook_path").split("/")
//val workspaceUrl=tagMap("browserHostName")
var workspaceUrl=https://adb-7614304971745696.16.azuredatabricks.net/
    
val workspaceName=dbutils.notebook().getContext().notebookPath.get
val notebookURL = workspaceUrl+"/?o="+tagMap("orgId")
val user = tagMap("user")
val name = notebookPath(notebookPath.size-1)
val notebookInfo = Map("notebookURL" -> notebookURL,
"user" -> user,
"workspaceName" ->workspaceName,
"workspaceUrl" -> workspaceUrl,                       
"name" -> name,
"mounts" -> dbutils.fs.ls("/FileStore/tables").map(_.path),
"timestamp" -> System.currentTimeMillis)
val notebookInfoJson = scala.util.parsing.json.JSONObject(notebookInfo)

class CustomFilter extends PostProcessingFilter {
  def this(conf: Configuration) = this()

  override def processExecutionEvent(event: ExecutionEvent, ctx: HarvestingContext): ExecutionEvent =
    event.withAddedExtra(Map("foo" -> "bar"))

  override def processExecutionPlan(plan: ExecutionPlan, ctx: HarvestingContext ): ExecutionPlan =
    plan.withAddedExtra(Map( "notebookInfo" -> notebookInfoJson))

  override def processReadOperation(op: ReadOperation, ctx: HarvestingContext ): ReadOperation =
    op.withAddedExtra(Map("foo" -> "bar"))

  override def processWriteOperation(op: WriteOperation, ctx: HarvestingContext): WriteOperation =
    op.withAddedExtra(Map("foo" -> "bar"))

  override def processDataOperation(op: DataOperation, ctx: HarvestingContext  ): DataOperation =
    op.withAddedExtra(Map("foo" -> "bar"))
}

val myInstance = new CustomFilter()

spark.enableLineageTracking(
  AgentConfig.builder()
    .postProcessingFilter(myInstance)
    .build()
)