Merge pull request #66 from civitaspo/develop

civitaspo · web-flow · commit c280ef642e2d · 2019-07-29T13:18:35.000+09:00
v0.2.4
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,3 +1,9 @@
+0.2.4 (2019-07-29)
+==================
+
+* [New Feature] Add `athena.table_exists?>` operator
+* [New Feature] Add `athena.partition_exists?>` operator
+
 0.2.3 (2019-07-19)
 ==================
 * [New Feature] Add `athena.drop_table_multi>` operator
diff --git a/README.md b/README.md
@@ -15,7 +15,7 @@ _export:
     repositories:
       - https://jitpack.io
     dependencies:
-      - pro.civitaspo:digdag-operator-athena:0.2.3
+      - pro.civitaspo:digdag-operator-athena:0.2.4
   athena:
     auth_method: profile
 
@@ -32,6 +32,8 @@ _export:
   output: s3://mybucket/prefix/
 ```
 
+See [examples](./example/example.dig) for more cases.
+
 # Configuration
 
 ## Remarks
@@ -112,6 +114,22 @@ Nothing
 
 Nothing
 
+## Configuration for `athena.partition_exists?>` operator
+
+### Options
+
+- **database**: The name of the database. (string, required)
+- **table**: The name of the partitioned table. (string, required)
+- **partition_kv**: key-value pairs for the partition (string to string map, required)
+- **with_location**: Check the partition existence with the location (boolean, default: `false`)
+- **error_if_not_exist**: Raise the exception if the partition or location does not exist. (boolean, default: `false`)
+- **catalog_id**: glue data catalog id if you use a catalog different from account/region default catalog. (string, optional)
+
+### Output Parameters
+
+- **athena.last_partition_exists.partition_exists**: `true` if the partition exists, or `false` (boolean)
+- **athena.last_partition_exists.location_exists**: `true` if the partition location exists, or `false`. `null` if not set **with_location** option is `true`. (boolean)
+
 ## Configuration for `athena.apas>` operator
 
 `apas` means *Add a partition as select* that creates a partition the query result is stored.
@@ -249,6 +267,22 @@ Nothing
 
 Nothing
 
+## Configuration for `athena.table_exists?>` operator
+
+### Options
+
+- **database**: The name of the database. (string, required)
+- **table**: The name of the table. (string, required)
+- **with_location**: Check the partition existence with the location (boolean, default: `false`)
+- **error_if_not_exist**: Raise the exception if the table or location does not exist. (boolean, default: `false`)
+- **catalog_id**: glue data catalog id if you use a catalog different from account/region default catalog. (string, optional)
+
+### Output Parameters
+
+- **athena.last_partition_exists.table_exists**: `true` if the table exists, or `false` (boolean)
+- **athena.last_partition_exists.location_exists**: `true` if the table location exists, or `false`. `null` if not set **with_location** option is `true`. (boolean)
+
+
 # Development
 
 ## Run an Example
diff --git a/build.gradle b/build.gradle
@@ -5,7 +5,7 @@ plugins {
 }
 
 group = 'pro.civitaspo'
-version = '0.2.3'
+version = '0.2.4'
 
 def digdagVersion = '0.9.37'
 def awsSdkVersion = "1.11.587"
diff --git a/example/example.dig b/example/example.dig
@@ -4,7 +4,7 @@ _export:
       - file://${repos}
       # - https://jitpack.io
     dependencies:
-      - pro.civitaspo:digdag-operator-athena:0.2.3
+      - pro.civitaspo:digdag-operator-athena:0.2.4
   athena:
     auth_method: profile
     value: 5
@@ -92,6 +92,27 @@ _export:
   save_mode: overwrite
 
 +step14:
+  +step1:
+    athena.partition_exists?>:
+    database: ${database}
+    table: hoge
+    with_location: true
+    partition_kv:
+      b: "9"
+      c: "10"
+  +step2:
+    echo>: ${athena.last_partition_exists}
+  +step3:
+    athena.partition_exists?>:
+    database: ${database}
+    table: hoge
+    partition_kv:
+      b: "9"
+      c: "10"
+  +step4:
+    echo>: ${athena.last_partition_exists}
+
++step15:
   loop>: 10
   _parallel: true
   _do:
@@ -101,15 +122,66 @@ _export:
     location: ${output}/hoge_${i}/
     save_mode: overwrite
 
-+step15:
++step16:
+  +step1:
+    athena.table_exists?>:
+    database: ${database}
+    table: hoge_1
+    with_location: true
+  +step2:
+    echo>: ${athena.last_table_exists}
+  +step3:
+    athena.table_exists?>:
+    database: ${database}
+    table: hoge
+  +step4:
+    echo>: ${athena.last_table_exists}
+
++step17:
   athena.drop_table_multi>:
   database: ${database}
   regexp: 'hoge_\d+'
   with_location: true
   limit: 3
 
-+step16:
++step18:
   athena.drop_table_multi>:
   database: ${database}
   regexp: 'hoge_\d+'
   with_location: true
+
++step19:
+  +step1:
+    athena.table_exists?>:
+    database: ${database}
+    table: hoge_1
+    with_location: true
+  +step2:
+    echo>: ${athena.last_table_exists}
+  +step3:
+    athena.table_exists?>:
+    database: ${database}
+    table: hoge
+  +step4:
+    echo>: ${athena.last_table_exists}
+
++step20:
+  +step1:
+    athena.partition_exists?>:
+    database: ${database}
+    table: hoge_1
+    with_location: true
+    partition_kv:
+      b: "9"
+      c: "10"
+  +step2:
+    echo>: ${athena.last_partition_exists}
+  +step3:
+    athena.partition_exists?>:
+    database: ${database}
+    table: hoge_1
+    partition_kv:
+      b: "9"
+      c: "10"
+  +step4:
+    echo>: ${athena.last_partition_exists}
diff --git a/src/main/scala/pro/civitaspo/digdag/plugin/athena/AthenaPlugin.scala b/src/main/scala/pro/civitaspo/digdag/plugin/athena/AthenaPlugin.scala
@@ -13,8 +13,10 @@ import pro.civitaspo.digdag.plugin.athena.ctas.AthenaCtasOperator
 import pro.civitaspo.digdag.plugin.athena.drop_partition.AthenaDropPartitionOperator
 import pro.civitaspo.digdag.plugin.athena.drop_table.AthenaDropTableOperator
 import pro.civitaspo.digdag.plugin.athena.drop_table_multi.AthenaDropTableMultiOperator
+import pro.civitaspo.digdag.plugin.athena.partition_exists.AthenaPartitionExistsOperator
 import pro.civitaspo.digdag.plugin.athena.preview.AthenaPreviewOperator
 import pro.civitaspo.digdag.plugin.athena.query.AthenaQueryOperator
+import pro.civitaspo.digdag.plugin.athena.table_exists.AthenaTableExistsOperator
 
 
 object AthenaPlugin
@@ -38,7 +40,9 @@ object AthenaPlugin
                 operatorFactory("athena.query", classOf[AthenaQueryOperator]),
                 operatorFactory("athena.preview", classOf[AthenaPreviewOperator]),
                 operatorFactory("athena.drop_table", classOf[AthenaDropTableOperator]),
-                operatorFactory("athena.drop_table_multi", classOf[AthenaDropTableMultiOperator])
+                operatorFactory("athena.drop_table_multi", classOf[AthenaDropTableMultiOperator]),
+                operatorFactory("athena.partition_exists?", classOf[AthenaPartitionExistsOperator]),
+                operatorFactory("athena.table_exists?", classOf[AthenaTableExistsOperator])
                 )
         }
 
diff --git a/src/main/scala/pro/civitaspo/digdag/plugin/athena/partition_exists/AthenaPartitionExistsOperator.scala b/src/main/scala/pro/civitaspo/digdag/plugin/athena/partition_exists/AthenaPartitionExistsOperator.scala
@@ -0,0 +1,73 @@
+package pro.civitaspo.digdag.plugin.athena.partition_exists
+
+
+import com.amazonaws.services.glue.model.Partition
+import com.google.common.collect.ImmutableList
+import io.digdag.client.config.{Config, ConfigKey}
+import io.digdag.spi.{OperatorContext, TaskResult, TemplateEngine}
+import pro.civitaspo.digdag.plugin.athena.AbstractAthenaOperator
+
+import scala.jdk.CollectionConverters._
+import scala.util.{Failure, Success, Try}
+
+
+class AthenaPartitionExistsOperator(operatorName: String,
+                                    context: OperatorContext,
+                                    systemConfig: Config,
+                                    templateEngine: TemplateEngine)
+    extends AbstractAthenaOperator(operatorName, context, systemConfig, templateEngine)
+{
+    protected val database: String = params.get("database", classOf[String])
+    protected val table: String = params.get("table", classOf[String])
+    protected val partitionKv: Map[String, String] = params.getMap("partition_kv", classOf[String], classOf[String]).asScala.toMap
+    protected val withLocation: Boolean = params.get("with_location", classOf[Boolean], false)
+    protected val errorIfNotExist: Boolean = params.get("error_if_not_exist", classOf[Boolean], false)
+    protected val catalogId: Option[String] = Option(params.getOptional("catalog_id", classOf[String]).orNull())
+
+
+    override def runTask(): TaskResult =
+    {
+        val storeParams: Config = cf.create()
+        val lastPartitionExistsParams = storeParams.getNestedOrSetEmpty("athena").getNestedOrSetEmpty("last_partition_exists")
+
+        run(storeParams = lastPartitionExistsParams)
+
+        val builder = TaskResult.defaultBuilder(request)
+        builder.resetStoreParams(ImmutableList.of(ConfigKey.of("athena", "last_partition_exists")))
+        builder.storeParams(storeParams)
+        builder.build()
+    }
+
+    protected def run(storeParams: Config): Unit =
+    {
+        if (aws.glue.partition.exists(catalogId, database, table, partitionKv)) {
+            storeParams.set("partition_exists", true)
+        }
+        else {
+            if (errorIfNotExist) throw new IllegalStateException(s"The partition{${partitionKv.mkString(",")}} does not exist.")
+            storeParams.set("partition_exists", false)
+        }
+
+        if (withLocation) {
+            Try(aws.glue.partition.describe(catalogId, database, table, partitionKv)) match {
+                case Success(p) =>
+                    val location: String = {
+                        val l = p.getStorageDescriptor.getLocation
+                        if (l.endsWith("/")) l
+                        else l + "/"
+                    }
+                    if (aws.s3.hasObjects(location)) {
+                        storeParams.set("location_exists", true)
+                    }
+                    else {
+                        if (errorIfNotExist) throw new IllegalStateException(s"The location '$location' does not exist, although the partition{${partitionKv.mkString(",")}} exists.")
+                        storeParams.set("location_exists", false)
+                    }
+
+                case Failure(e) =>
+                    logger.warn(s"The partition{${partitionKv.mkString(",")}} may not exist.", e)
+                    storeParams.set("location_exists", false)
+            }
+        }
+    }
+}
diff --git a/src/main/scala/pro/civitaspo/digdag/plugin/athena/table_exists/AthenaTableExistsOperator.scala b/src/main/scala/pro/civitaspo/digdag/plugin/athena/table_exists/AthenaTableExistsOperator.scala
@@ -0,0 +1,70 @@
+package pro.civitaspo.digdag.plugin.athena.table_exists
+
+
+import com.google.common.collect.ImmutableList
+import io.digdag.client.config.{Config, ConfigKey}
+import io.digdag.spi.{OperatorContext, TaskResult, TemplateEngine}
+import pro.civitaspo.digdag.plugin.athena.AbstractAthenaOperator
+
+import scala.util.{Failure, Success, Try}
+
+
+class AthenaTableExistsOperator(operatorName: String,
+                                context: OperatorContext,
+                                systemConfig: Config,
+                                templateEngine: TemplateEngine)
+    extends AbstractAthenaOperator(operatorName, context, systemConfig, templateEngine)
+{
+
+    val database: String = params.get("database", classOf[String])
+    val table: String = params.get("table", classOf[String])
+    val withLocation: Boolean = params.get("with_location", classOf[Boolean], false)
+    val errorIfNotExist: Boolean = params.get("error_if_not_exist", classOf[Boolean], false)
+    val catalogId: Option[String] = Option(params.getOptional("catalog_id", classOf[String]).orNull())
+
+    override def runTask(): TaskResult =
+    {
+        val storeParams: Config = cf.create()
+        val lastTableExistsParams = storeParams.getNestedOrSetEmpty("athena").getNestedOrSetEmpty("last_table_exists")
+
+        run(storeParams = lastTableExistsParams)
+
+        val builder = TaskResult.defaultBuilder(request)
+        builder.resetStoreParams(ImmutableList.of(ConfigKey.of("athena", "last_table_exists")))
+        builder.storeParams(storeParams)
+        builder.build()
+    }
+
+    protected def run(storeParams: Config): Unit =
+    {
+        if (aws.glue.table.exists(catalogId, database, table)) {
+            storeParams.set("table_exists", true)
+        }
+        else {
+            if (errorIfNotExist) throw new IllegalStateException(s"The table '$database.$table' does not exist.")
+            storeParams.set("table_exists", false)
+        }
+
+        if (withLocation) {
+            Try(aws.glue.table.describe(catalogId, database, table)) match {
+                case Success(t) =>
+                    val location: String = {
+                        val l = t.getStorageDescriptor.getLocation
+                        if (l.endsWith("/")) l
+                        else l + "/"
+                    }
+                    if (aws.s3.hasObjects(location)) {
+                        storeParams.set("location_exists", true)
+                    }
+                    else {
+                        if (errorIfNotExist) throw new IllegalStateException(s"The location '$location' does not exist, although the table '$database.$table' exists.")
+                        storeParams.set("location_exists", false)
+                    }
+
+                case Failure(e) =>
+                    logger.warn(s"The table '$database.$table' may not exist.", e)
+                    storeParams.set("location_exists", false)
+            }
+        }
+    }
+}

Original file line number	Diff line number	Diff line change
`@@ -5,7 +5,7 @@ plugins {`
`5`	`5`	`}`
`6`	`6`
`7`	`7`	`group = 'pro.civitaspo'`
`8`		`-version = '0.2.3'`
	`8`	`+version = '0.2.4'`
`9`	`9`
`10`	`10`	`def digdagVersion = '0.9.37'`
`11`	`11`	`def awsSdkVersion = "1.11.587"`