1 /***********************************************************************
2  * Copyright (c) 2013-2025 Commonwealth Computer Research, Inc.
3  * All rights reserved. This program and the accompanying materials
4  * are made available under the terms of the Apache License, Version 2.0
5  * which accompanies this distribution and is available at
6  * http://www.opensource.org/licenses/apache2.0.php.
7  ***********************************************************************/
8 
9 package org.locationtech.geomesa.accumulo.data.writer
10 package tx
11 
12 import org.apache.accumulo.core.client.{ConditionalWriter, ConditionalWriterConfig, IsolatedScanner}
13 import org.apache.accumulo.core.conf.ClientProperty
14 import org.apache.accumulo.core.data.{Key, PartialKey}
15 import org.apache.hadoop.io.Text
16 import org.geotools.api.feature.simple.SimpleFeatureType
17 import org.locationtech.geomesa.accumulo.data.AccumuloDataStore
18 import org.locationtech.geomesa.accumulo.data.writer.tx.ConditionalWriteException.ConditionalWriteStatus
19 import org.locationtech.geomesa.accumulo.data.writer.tx.MutationBuilder.{AppendBuilder, DeleteBuilder, UpdateBuilder}
20 import org.locationtech.geomesa.index.api.IndexAdapter.BaseIndexWriter
21 import org.locationtech.geomesa.index.api.WritableFeature.FeatureWrapper
22 import org.locationtech.geomesa.index.api._
23 import org.locationtech.geomesa.utils.io.{CloseQuietly, WithClose}
24 
25 import java.util.concurrent.TimeUnit
26 import scala.collection.mutable.ArrayBuffer
27 
28 /**
29  * Accumulo atomic index writer implementation
30  *
31  * @param ds data store
32  * @param sft simple feature type
33  * @param indices indices to write to
34  * @param wrapper feature wrapper
35  * @param partition partition to write to (if partitioned schema)
36  */
37 class AccumuloAtomicIndexWriter(
38     ds: AccumuloDataStore,
39     sft: SimpleFeatureType,
40     indices: Seq[GeoMesaFeatureIndex[_, _]],
41     wrapper: FeatureWrapper[WritableFeature],
42     partition: Option[String]
43   ) extends BaseIndexWriter[WritableFeature](indices, wrapper) {
44 
45   import scala.collection.JavaConverters._
46 
47   // should always be writing to a single table here
48   private val tables = indices.map(_.getTableName(partition))
49 
50   private val writers: Array[ConditionalWriter] = {
51     val config = new ConditionalWriterConfig()
52     config.setAuthorizations(ds.auths)
53     val maxThreads = ClientProperty.BATCH_WRITER_THREADS_MAX.getInteger(ds.connector.properties())
54     if (maxThreads != null) {
55       config.setMaxWriteThreads(maxThreads)
56     }
57     val timeout = ClientProperty.BATCH_WRITER_TIMEOUT_MAX.getTimeInMillis(ds.connector.properties())
58     if (timeout != null) {
59       config.setTimeout(timeout, TimeUnit.MILLISECONDS)
60     }
61     tables.map(ds.connector.createConditionalWriter(_, config)).toArray
62   }
63 
64   private val colFamilyMappings = indices.map(ColumnFamilyMapper.apply).toArray
65   private val visCache = new VisibilityCache()
66 
67   override protected def append(feature: WritableFeature, values: Array[RowKeyValue[_]]): Unit =
68     mutate(feature.feature.getID, buildMutations(values, AppendBuilder.apply))
69 
70   override protected def delete(feature: WritableFeature, values: Array[RowKeyValue[_]]): Unit =
71     mutate(feature.feature.getID, buildMutations(values, DeleteBuilder.apply))
72 
73   override protected def update(
74       feature: WritableFeature,
75       values: Array[RowKeyValue[_]],
76       previous: WritableFeature,
77       previousValues: Array[RowKeyValue[_]]): Unit = {
78     val mutations = Array.ofDim[Seq[MutationBuilder]](values.length)
79     // note: these are temporary conditions - we update them below
80     val updates = buildMutations(values, AppendBuilder.apply)
81     val deletes = buildMutations(previousValues, DeleteBuilder.apply)
82     var i = 0
83     while (i < values.length) {
84       val dels = ArrayBuffer(deletes(i): _*)
85       val puts = updates(i).map { mutations =>
86         // find any previous values that will be updated
87         val d = dels.indexWhere(p => java.util.Arrays.equals(p.row, mutations.row))
88         if (d == -1) { mutations } else {
89           // note: side-effect
90           val toUpdate = dels.remove(d)
91           // any previous values that are overwritten are added as conditions on the put
92           // any previous values that aren't overwritten need to be deleted
93           val remaining = toUpdate.kvs.filterNot(kv => mutations.kvs.exists(_.equalKey(kv)))
94           if (remaining.nonEmpty) {
95             dels.append(toUpdate.copy(kvs = remaining))
96           }
97           UpdateBuilder(mutations.row, mutations.kvs, toUpdate.kvs)
98         }
99       }
100       // note: order here is important, we need puts to operate before deletes in order for conditions to succeed
101       mutations(i) = puts ++ dels
102       i += 1
103     }
104     mutate(feature.feature.getID, mutations)
105   }
106 
107   /**
108    * Apply the mutations
109    *
110    * @param id feature id
111    * @param mutations mutations to apply
112    */
113   @throws[ConditionalWriteException]
114   private def mutate(id: String, mutations: Array[Seq[MutationBuilder]]): Unit = {
115     val errors = ArrayBuffer.empty[ConditionalWriteStatus]
116     val successes = ArrayBuffer.empty[(Int, MutationBuilder)]
117     var i = 0
118     while (i < mutations.length) {
119       // note: mutations need to be applied in order so that conditions are correct
120       mutations(i).foreach { m =>
121         val status = writers(i).write(m.apply()).getStatus
122         if (status == ConditionalWriter.Status.ACCEPTED ||
123             (status == ConditionalWriter.Status.UNKNOWN && verifyWrite(tables(i), m))) {
124           successes += i -> m
125         } else {
126           errors += ConditionalWriteStatus(indices(i), m.name, status)
127         }
128       }
129       i += 1
130     }
131     if (errors.nonEmpty) {
132       // revert any writes we made
133       successes.foreach { case (i, m) =>
134         // note: if these are rejected, some other update has come through and we don't want to overwrite it
135         writers(i).write(m.invert())
136       }
137       throw ConditionalWriteException(id, errors.toSeq)
138     }
139   }
140 
141   /**
142    * Verify if a mutation was successfully applied or not
143    *
144    * @param table table being mutated
145    * @param mutation mutation being applied
146    * @return true if the current state of the table reflects the mutation's intended result
147    */
148   private def verifyWrite(table: String, mutation: MutationBuilder): Boolean = {
149     WithClose(new IsolatedScanner(ds.connector.createScanner(table, ds.auths))) { scanner =>
150       val key = new Key(mutation.row)
151       scanner.setRange(new org.apache.accumulo.core.data.Range(key, true, key.followingKey(PartialKey.ROW), false))
152       val found = scanner.iterator().asScala.toList
153       mutation.apply().getUpdates.asScala.forall { update =>
154         if (update.isDeleted) {
155           found.forall { entry =>
156             !compare(entry.getKey.getColumnFamily, update.getColumnFamily) ||
157                 !compare(entry.getKey.getColumnQualifier, update.getColumnQualifier) ||
158                 !compare(entry.getKey.getColumnVisibility, update.getColumnVisibility) ||
159                 entry.getKey.isDeleted
160           }
161         } else {
162           found.exists { entry =>
163             compare(entry.getKey.getColumnFamily, update.getColumnFamily) &&
164                 compare(entry.getKey.getColumnQualifier, update.getColumnQualifier) &&
165                 compare(entry.getKey.getColumnVisibility, update.getColumnVisibility) &&
166                 entry.getValue.compareTo(update.getValue) == 0
167           }
168         }
169       }
170     }
171   }
172 
173   private def buildMutations[T <: MutationBuilder](
174       values: Array[RowKeyValue[_]],
175       builderFactory: (Array[Byte], Seq[MutationValue]) => T): Array[Seq[T]] = {
176     val mutations = Array.ofDim[Seq[T]](values.length)
177     var i = 0
178     while (i < values.length) {
179       mutations(i) = values(i) match {
180         case kv: SingleRowKeyValue[_] =>
181           val values = kv.values.map { v =>
182             MutationValue(colFamilyMappings(i)(v.cf), v.cq, visCache(v.vis), v.value)
183           }
184           Seq(builderFactory(kv.row, values))
185 
186         case mkv: MultiRowKeyValue[_] =>
187           mkv.rows.map { row =>
188             val values = mkv.values.map { v =>
189               MutationValue(colFamilyMappings(i)(v.cf), v.cq, visCache(v.vis), v.value)
190             }
191             builderFactory(row, values)
192           }
193       }
194       i += 1
195     }
196     mutations
197   }
198 
199   private def compare(text: Text, bytes: Array[Byte]): Boolean = text.compareTo(bytes, 0, bytes.length) == 0
200 
201   override def flush(): Unit = {} // there is no batching here, every single write gets flushed
202 
203   override def close(): Unit = CloseQuietly(writers).foreach(e => throw e)
204 }