chore: add Micronaut Data layer for Organization Payment Config (#13651)

airbyte-config/config-models/src/main/resources/types/OrganizationPaymentConfig.yaml

@@ -0,0 +1,38 @@
+---
+"$schema": http://json-schema.org/draft-07/schema#
+"$id": https://github.com/airbytehq/airbyte/blob/master/airbyte-config/models/src/main/resources/types/OrganizationPaymentConfig.yaml
+title: OrganizationPaymentConfig
+description: Organization Payment Config
+type: object
+required:
+  - organizationId
+  - payment_status
+  - created_at
+  - updated_at
+additionalProperties: true
+properties:
+  organizationId:
+    description: ID of the associated organization
+    type: string
+    format: uuid
+  payment_provider_id:
+    description: ID of the external payment provider (ex. a Stripe Customer ID)
+    type: string
+  payment_status:
+    description: Payment status for the organization
+    $ref: PaymentStatus.yaml
+  grace_period_end_at:
+    description: If set, the date at which the organization's grace period ends and syncs will be disabled
+    type: integer
+    format: int64
+  usage_category_override:
+    description: If set, the usage category that the organization should always be billed with
+    $ref: UsageCategoryOverride.yaml
+  created_at:
+    description: Creation timestamp of the organization payment config
+    type: integer
+    format: int64
+  updated_at:
+    description: Last updated timestamp of the organization payment config
+    type: integer
+    format: int64

airbyte-config/config-models/src/main/resources/types/PaymentStatus.yaml

@@ -0,0 +1,13 @@
+---
+"$schema": http://json-schema.org/draft-07/schema#
+"$id": https://github.com/airbytehq/airbyte-platform/blob/main/airbyte-config/config-models/src/main/resources/types/PaymentStatus.yaml
+title: PaymentStatus
+description: Payment Status for an Organization Payment Config
+type: string
+enum:
+  - uninitialized
+  - okay
+  - grace_period
+  - disabled
+  - locked
+  - manual

airbyte-config/config-models/src/main/resources/types/UsageCategoryOverride.yaml

@@ -0,0 +1,9 @@
+---
+"$schema": http://json-schema.org/draft-07/schema#
+"$id": https://github.com/airbytehq/airbyte-platform/blob/main/airbyte-config/config-models/src/main/resources/types/UsageCategoryOverride.yaml
+title: UsageCategoryOverride
+description: Usage Category Override for an Organization Payment Config
+type: string
+enum:
+  - free
+  - internal

airbyte-data/src/main/java/io/airbyte/data/services/impls/jooq/OrganizationServiceJooqImpl.java

@@ -54,7 +54,7 @@ public Optional<Organization> getOrganization(final UUID organizationId) throws
   }
 
   @Override
-  public Optional<Organization> getOrganizationForWorkspaceId(UUID workspaceId) throws IOException {
+  public Optional<Organization> getOrganizationForWorkspaceId(UUID workspaceId) {
     throw new UnsupportedOperationException("Not implemented - use OrganizationServiceDataImpl instead");
   }
 

airbyte-data/src/main/kotlin/io/airbyte/data/repositories/OrganizationPaymentConfigRepository.kt

@@ -0,0 +1,10 @@
+package io.airbyte.data.repositories
+
+import io.airbyte.data.repositories.entities.OrganizationPaymentConfig
+import io.micronaut.data.jdbc.annotation.JdbcRepository
+import io.micronaut.data.model.query.builder.sql.Dialect
+import io.micronaut.data.repository.PageableRepository
+import java.util.UUID
+
+@JdbcRepository(dialect = Dialect.POSTGRES, dataSource = "config")
+interface OrganizationPaymentConfigRepository : PageableRepository<OrganizationPaymentConfig, UUID>

airbyte-data/src/main/kotlin/io/airbyte/data/repositories/entities/OrganizationPaymentConfig.kt

@@ -0,0 +1,27 @@
+package io.airbyte.data.repositories.entities
+
+import io.airbyte.db.instance.configs.jooq.generated.enums.PaymentStatus
+import io.airbyte.db.instance.configs.jooq.generated.enums.UsageCategoryOverride
+import io.micronaut.data.annotation.DateCreated
+import io.micronaut.data.annotation.DateUpdated
+import io.micronaut.data.annotation.Id
+import io.micronaut.data.annotation.MappedEntity
+import io.micronaut.data.annotation.TypeDef
+import io.micronaut.data.model.DataType
+import java.util.UUID
+
+@MappedEntity("organization_payment_config")
+open class OrganizationPaymentConfig(
+  @field:Id
+  var organizationId: UUID,
+  var paymentProviderId: String? = null,
+  @field:TypeDef(type = DataType.OBJECT)
+  var paymentStatus: PaymentStatus = PaymentStatus.uninitialized,
+  var gracePeriodEndAt: java.time.OffsetDateTime? = null,
+  @field:TypeDef(type = DataType.OBJECT)
+  var usageCategoryOverride: UsageCategoryOverride? = null,
+  @DateCreated
+  var createdAt: java.time.OffsetDateTime? = null,
+  @DateUpdated
+  var updatedAt: java.time.OffsetDateTime? = null,
+)

airbyte-data/src/main/kotlin/io/airbyte/data/services/OrganizationPaymentConfigService.kt

@@ -0,0 +1,8 @@
+package io.airbyte.data.services
+
+import io.airbyte.config.OrganizationPaymentConfig
+import java.util.UUID
+
+interface OrganizationPaymentConfigService {
+  fun findByOrganizationId(organizationId: UUID): OrganizationPaymentConfig?
+}

airbyte-data/src/main/kotlin/io/airbyte/data/services/impls/data/OrganizationPaymentConfigServiceDataImpl.kt

@@ -0,0 +1,16 @@
+package io.airbyte.data.services.impls.data
+
+import io.airbyte.config.OrganizationPaymentConfig
+import io.airbyte.data.repositories.OrganizationPaymentConfigRepository
+import io.airbyte.data.services.OrganizationPaymentConfigService
+import io.airbyte.data.services.impls.data.mappers.toConfigModel
+import jakarta.inject.Singleton
+import java.util.UUID
+
+@Singleton
+class OrganizationPaymentConfigServiceDataImpl(
+  private val organizationPaymentConfigRepository: OrganizationPaymentConfigRepository,
+) : OrganizationPaymentConfigService {
+  override fun findByOrganizationId(organizationId: UUID): OrganizationPaymentConfig? =
+    organizationPaymentConfigRepository.findById(organizationId).orElse(null)?.toConfigModel()
+}

airbyte-data/src/main/kotlin/io/airbyte/data/services/impls/data/mappers/OrganizationPaymentConfigMapper.kt

@@ -0,0 +1,67 @@
+package io.airbyte.data.services.impls.data.mappers
+
+import io.airbyte.config.OrganizationPaymentConfig as ModelOrganizationPaymentConfig
+import io.airbyte.config.OrganizationPaymentConfig.PaymentStatus as ModelPaymentStatus
+import io.airbyte.config.OrganizationPaymentConfig.UsageCategoryOverride as ModelUsageCategoryOverride
+import io.airbyte.data.repositories.entities.OrganizationPaymentConfig as EntityOrganizationPaymentConfig
+import io.airbyte.db.instance.configs.jooq.generated.enums.PaymentStatus as EntityPaymentStatus
+import io.airbyte.db.instance.configs.jooq.generated.enums.UsageCategoryOverride as EntityUsageCategoryOverride
+
+fun EntityOrganizationPaymentConfig.toConfigModel(): ModelOrganizationPaymentConfig =
+  ModelOrganizationPaymentConfig()
+    .withOrganizationId(this.organizationId)
+    .withPaymentProviderId(this.paymentProviderId)
+    .withPaymentStatus(this.paymentStatus.toConfigModel())
+    .withGracePeriodEndAt(this.gracePeriodEndAt?.toEpochSecond())
+    .withUsageCategoryOverride(this.usageCategoryOverride?.toConfigModel())
+    .withCreatedAt(this.createdAt?.toEpochSecond())
+    .withUpdatedAt(this.updatedAt?.toEpochSecond())
+
+fun ModelOrganizationPaymentConfig.toEntity(): EntityOrganizationPaymentConfig =
+  EntityOrganizationPaymentConfig(
+    organizationId = this.organizationId,
+    paymentProviderId = this.paymentProviderId,
+    paymentStatus = this.paymentStatus.toEntity(),
+    gracePeriodEndAt =
+      this.gracePeriodEndAt?.let {
+        java.time.OffsetDateTime.ofInstant(
+          java.time.Instant.ofEpochSecond(it),
+          java.time.ZoneOffset.UTC,
+        )
+      },
+    usageCategoryOverride = this.usageCategoryOverride?.toEntity(),
+    createdAt = this.createdAt?.let { java.time.OffsetDateTime.ofInstant(java.time.Instant.ofEpochSecond(it), java.time.ZoneOffset.UTC) },
+    updatedAt = this.updatedAt?.let { java.time.OffsetDateTime.ofInstant(java.time.Instant.ofEpochSecond(it), java.time.ZoneOffset.UTC) },
+  )
+
+fun EntityPaymentStatus.toConfigModel(): ModelPaymentStatus =
+  when (this) {
+    EntityPaymentStatus.uninitialized -> ModelPaymentStatus.UNINITIALIZED
+    EntityPaymentStatus.okay -> ModelPaymentStatus.OKAY
+    EntityPaymentStatus.grace_period -> ModelPaymentStatus.GRACE_PERIOD
+    EntityPaymentStatus.disabled -> ModelPaymentStatus.DISABLED
+    EntityPaymentStatus.locked -> ModelPaymentStatus.LOCKED
+    EntityPaymentStatus.manual -> ModelPaymentStatus.MANUAL
+  }
+
+fun ModelPaymentStatus.toEntity(): EntityPaymentStatus =
+  when (this) {
+    ModelPaymentStatus.UNINITIALIZED -> EntityPaymentStatus.uninitialized
+    ModelPaymentStatus.OKAY -> EntityPaymentStatus.okay
+    ModelPaymentStatus.GRACE_PERIOD -> EntityPaymentStatus.grace_period
+    ModelPaymentStatus.DISABLED -> EntityPaymentStatus.disabled
+    ModelPaymentStatus.LOCKED -> EntityPaymentStatus.locked
+    ModelPaymentStatus.MANUAL -> EntityPaymentStatus.manual
+  }
+
+fun EntityUsageCategoryOverride.toConfigModel(): ModelUsageCategoryOverride =
+  when (this) {
+    EntityUsageCategoryOverride.free -> ModelUsageCategoryOverride.FREE
+    EntityUsageCategoryOverride.internal -> ModelUsageCategoryOverride.INTERNAL
+  }
+
+fun ModelUsageCategoryOverride.toEntity(): EntityUsageCategoryOverride =
+  when (this) {
+    ModelUsageCategoryOverride.FREE -> EntityUsageCategoryOverride.free
+    ModelUsageCategoryOverride.INTERNAL -> EntityUsageCategoryOverride.internal
+  }

airbyte-data/src/test/kotlin/io/airbyte/data/repositories/AbstractConfigRepositoryTest.kt

@@ -90,4 +90,5 @@ abstract class AbstractConfigRepositoryTest {
   val authRefreshTokenRepository = context.getBean(AuthRefreshTokenRepository::class.java)!!
   val organizationRepository = context.getBean(OrganizationRepository::class.java)!!
   val workspaceRepository = context.getBean(WorkspaceRepository::class.java)!!
+  val organizationPaymentConfigRepository = context.getBean(OrganizationPaymentConfigRepository::class.java)!!
 }

airbyte-data/src/test/kotlin/io/airbyte/data/repositories/OrganizationPaymentConfigRepositoryTest.kt

@@ -0,0 +1,52 @@
+package io.airbyte.data.repositories
+
+import io.airbyte.data.repositories.entities.Organization
+import io.airbyte.data.repositories.entities.OrganizationPaymentConfig
+import io.airbyte.db.instance.configs.jooq.generated.enums.PaymentStatus
+import io.micronaut.test.extensions.junit5.annotation.MicronautTest
+import org.junit.jupiter.api.AfterEach
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Assertions.assertNotNull
+import org.junit.jupiter.api.Assertions.assertNull
+import org.junit.jupiter.api.Test
+
+@MicronautTest
+internal class OrganizationPaymentConfigRepositoryTest : AbstractConfigRepositoryTest() {
+  @AfterEach
+  fun tearDown() {
+    organizationPaymentConfigRepository.deleteAll()
+    organizationRepository.deleteAll()
+  }
+
+  @Test
+  fun `test db insertion and retrieval`() {
+    val organization =
+      Organization(
+        name = "Airbyte Inc.",
+        email = "[email protected]",
+      )
+    val persistedOrg = organizationRepository.save(organization)
+
+    val paymentConfig =
+      OrganizationPaymentConfig(
+        organizationId = persistedOrg.id!!,
+      )
+
+    val countBeforeSave = organizationPaymentConfigRepository.count()
+    assertEquals(0L, countBeforeSave)
+
+    organizationPaymentConfigRepository.save(paymentConfig)
+
+    val countAfterSave = organizationPaymentConfigRepository.count()
+    assertEquals(1L, countAfterSave)
+
+    val persistedPaymentConfig = organizationPaymentConfigRepository.findById(persistedOrg.id).get()
+    assertEquals(persistedOrg.id, persistedPaymentConfig.organizationId)
+    assertNull(persistedPaymentConfig.paymentProviderId)
+    assertEquals(PaymentStatus.uninitialized, persistedPaymentConfig.paymentStatus)
+    assertNull(persistedPaymentConfig.gracePeriodEndAt)
+    assertNull(persistedPaymentConfig.usageCategoryOverride)
+    assertNotNull(persistedPaymentConfig.createdAt)
+    assertNotNull(persistedPaymentConfig.updatedAt)
+  }
+}

airbyte-data/src/test/kotlin/io/airbyte/data/services/impls/data/OrganizationPaymentConfigServiceDataImplTest.kt

@@ -0,0 +1,66 @@
+package io.airbyte.data.services.impls.data
+
+import io.airbyte.data.repositories.OrganizationPaymentConfigRepository
+import io.airbyte.data.repositories.entities.OrganizationPaymentConfig
+import io.airbyte.db.instance.configs.jooq.generated.enums.PaymentStatus
+import io.airbyte.db.instance.configs.jooq.generated.enums.UsageCategoryOverride
+import io.mockk.clearAllMocks
+import io.mockk.every
+import io.mockk.mockk
+import io.mockk.verify
+import org.junit.jupiter.api.AfterEach
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Assertions.assertNotNull
+import org.junit.jupiter.api.Assertions.assertNull
+import org.junit.jupiter.api.BeforeEach
+import org.junit.jupiter.api.Test
+import java.time.OffsetDateTime
+import java.util.Optional
+import java.util.UUID
+
+internal class OrganizationPaymentConfigServiceDataImplTest {
+  private val organizationPaymentConfigRepository = mockk<OrganizationPaymentConfigRepository>()
+  private val organizationPaymentConfigService = OrganizationPaymentConfigServiceDataImpl(organizationPaymentConfigRepository)
+
+  private lateinit var testOrganizationId: UUID
+
+  @BeforeEach
+  fun setup() {
+    clearAllMocks()
+    testOrganizationId = UUID.randomUUID()
+    val orgPaymentConfig =
+      OrganizationPaymentConfig(
+        organizationId = testOrganizationId,
+        paymentProviderId = "provider-id",
+        paymentStatus = PaymentStatus.grace_period,
+        gracePeriodEndAt = OffsetDateTime.now().plusDays(30),
+        usageCategoryOverride = UsageCategoryOverride.internal,
+        createdAt = OffsetDateTime.now(),
+        updatedAt = OffsetDateTime.now(),
+      )
+    every { organizationPaymentConfigRepository.findById(testOrganizationId) } returns Optional.of(orgPaymentConfig)
+    every { organizationPaymentConfigRepository.findById(not(testOrganizationId)) } returns Optional.empty()
+  }
+
+  @AfterEach
+  fun tearDown() {
+    clearAllMocks()
+  }
+
+  @Test
+  fun `test find by organization id`() {
+    val result = organizationPaymentConfigService.findByOrganizationId(testOrganizationId)
+    assertNotNull(result)
+    assertEquals(testOrganizationId, result?.organizationId)
+    assertEquals("provider-id", result?.paymentProviderId)
+    verify { organizationPaymentConfigRepository.findById(testOrganizationId) }
+  }
+
+  @Test
+  fun `test find by non-existent organization id`() {
+    val nonExistentId = UUID.randomUUID()
+    val result = organizationPaymentConfigService.findByOrganizationId(nonExistentId)
+    assertNull(result)
+    verify { organizationPaymentConfigRepository.findById(nonExistentId) }
+  }
+}

airbyte-data/src/test/resources/application-test.yaml

@@ -0,0 +1,4 @@
+# Uncomment to see generated SQL in test output
+# logger:
+#   levels:
+#     io.micronaut.data.query: TRACE