Add tasks based ingestion

druid.go

@@ -292,10 +292,6 @@ func (c *Client) Tasks() *TasksService {
 	return &TasksService{client: c}
 }
 
-func (c *Client) Metadata(options ...metadataOption) *MetadataService {
-	return NewMetadataService(c, options...)
-}
-
 func WithBasicAuth(username, password string) ClientOption {
 	return func(opts *clientOptions) {
 		opts.username = username

metadata.go → metadata_test.go

@@ -9,6 +9,10 @@ import (
 	"github.com/h2oai/go-druid/builder/query"
 )
 
+func (c *Client) metadata(options ...metadataOption) *metadataService {
+	return newMetadataService(c, options...)
+}
+
 type count struct {
 	Cnt int `json:"cnt"`
 }
@@ -20,22 +24,23 @@ type metadataOptions struct {
 
 type metadataOption func(*metadataOptions)
 
-// MetadataService is a service that runs druid metadata requests using druid SQL API.
-type MetadataService struct {
+// metadataService is a service that runs druid metadata requests using druid SQL API.
+// NOTE: for internal tests use only, not representing official druid API.
+type metadataService struct {
 	client         *Client
 	tickerDuration time.Duration
 	awaitTimeout   time.Duration
 }
 
-func NewMetadataService(client *Client, options ...metadataOption) *MetadataService {
+func newMetadataService(client *Client, options ...metadataOption) *metadataService {
 	opts := &metadataOptions{
 		tickerDuration: 500 * time.Millisecond,
 		awaitTimeout:   180 * time.Second,
 	}
 	for _, opt := range options {
 		opt(opts)
 	}
-	md := &MetadataService{
+	md := &metadataService{
 		client:         client,
 		tickerDuration: opts.tickerDuration,
 		awaitTimeout:   opts.awaitTimeout,
@@ -62,8 +67,9 @@ func fillDataSourceName(in string, ds string) string {
 	return strings.Replace(in, "${{ datasource }}", ds, 1)
 }
 
-// AwaitDataSourceAvailable awaits for a datasource to be visible in druid table listing.
-func (md *MetadataService) AwaitDataSourceAvailable(dataSourceName string) error {
+// awaitDataSourceAvailable awaits for a datasource to be visible in druid table listing.
+// NOTE: for internal tests use only, not representing official druid API.
+func (md *metadataService) awaitDataSourceAvailable(dataSourceName string) error {
 	ticker := time.NewTicker(md.tickerDuration)
 	defer ticker.Stop()
 	afterTimeout := time.After(md.awaitTimeout)
@@ -83,16 +89,17 @@ func (md *MetadataService) AwaitDataSourceAvailable(dataSourceName string) error
 				return nil
 			}
 		case <-afterTimeout:
-			return errors.New("AwaitDataSourceAvailable timeout")
+			return errors.New("awaitDataSourceAvailable timeout")
 		}
 	}
 }
 
 //go:embed sql/datasource_records.sql
 var datasourceRecordsQuery string
 
-// AwaitRecordsCount awaits for specific recordsCount in a given datasource.
-func (md *MetadataService) AwaitRecordsCount(dataSourceName string, recordsCount int) error {
+// awaitRecordsCount awaits for specific recordsCount in a given datasource.
+// NOTE: not safe and intended for internal tests use only. Not representing official druid API.
+func (md *metadataService) awaitRecordsCount(dataSourceName string, recordsCount int) error {
 	ticker := time.NewTicker(md.tickerDuration)
 	defer ticker.Stop()
 	q := query.NewSQL()
@@ -111,7 +118,7 @@ func (md *MetadataService) AwaitRecordsCount(dataSourceName string, recordsCount
 				return nil
 			}
 		case <-afterTimeout:
-			return errors.New("AwaitRecordsCount timeout")
+			return errors.New("awaitRecordsCount timeout")
 		}
 	}
 }

sql/datasource_available.sql

@@ -1,3 +1,3 @@
 SELECT count(*) cnt
 FROM INFORMATION_SCHEMA.TABLES
-WHERE TABLE_SCHEMA = 'druid' AND TABLE_NAME=?
+WHERE TABLE_SCHEMA = 'druid' AND TABLE_NAME=?

sql/datasource_records.sql

@@ -1,2 +1,2 @@
 SELECT count(*) cnt
-FROM "${{ datasource }}" ds
+FROM "${{ datasource }}" ds

task_types.go

@@ -43,10 +43,6 @@ func defaultTaskIngestionSpec() *TaskIngestionSpec {
 		Spec: &IngestionSpecData{
 			DataSchema: &DataSchema{
 				DataSource: "some_datasource",
-				TimeStampSpec: &TimestampSpec{
-					Column: "ts",
-					Format: "auto",
-				},
 				GranularitySpec: &GranularitySpec{
 					Type:               "uniform",
 					SegmentGranularity: "DAY",
@@ -63,16 +59,9 @@ func defaultTaskIngestionSpec() *TaskIngestionSpec {
 			IOConfig: &IOConfig{
 				Type: "index_parallel",
 				InputSource: &InputSource{
-					Type: "sql",
-					Database: &Database{
-						Type: "postgresql",
-						ConnectorConfig: &ConnectorConfig{
-							ConnectURI: "jdbc:postgresql://host:port/schema",
-							User:       "user",
-							Password:   "password",
-						},
-					},
-					SQLs: []string{},
+					Type:     "sql",
+					Database: &Database{},
+					SQLs:     []string{},
 				},
 				InputFormat: &InputFormat{
 					Type: "json",
@@ -98,6 +87,18 @@ func SetTaskType(stype string) TaskIngestionSpecOptions {
 	}
 }
 
+// SetTaskTimestampColumn sets the type of the task IOConfig.
+func SetTaskTimestampColumn(column string) TaskIngestionSpecOptions {
+	return func(spec *TaskIngestionSpec) {
+		if column != "" {
+			spec.Spec.DataSchema.TimeStampSpec = &TimestampSpec{
+				Column: column,
+				Format: "auto",
+			}
+		}
+	}
+}
+
 // SetTaskDataSource sets the destination datasource of the task IOConfig.
 func SetTaskDataSource(datasource string) TaskIngestionSpecOptions {
 	return func(spec *TaskIngestionSpec) {

tasks_test.go

@@ -35,8 +35,8 @@ var testObjects = []testDO{
 	},
 }
 
-// TriggerIngestionTask initiates inline ingestion task with druid client.
-func TriggerIngestionTask[T any](d *Client, dataSourceName string, entries []T) (string, error) {
+// triggerIngestionTask initiates inline ingestion task with druid client.
+func triggerIngestionTask[T any](d *Client, dataSourceName string, entries []T) (string, error) {
 	csvEntriesBuff := &bytes.Buffer{}
 
 	err := gocsv.MarshalWithoutHeaders(entries, csvEntriesBuff)
@@ -46,6 +46,7 @@ func TriggerIngestionTask[T any](d *Client, dataSourceName string, entries []T)
 
 	var spec = NewTaskIngestionSpec(
 		SetTaskType("index_parallel"),
+		SetTaskTimestampColumn("ts"),
 		SetTaskDataSource(dataSourceName),
 		SetTaskTuningConfig("index_parallel", 25000, 5000000),
 		SetTaskIOConfigType("index_parallel"),
@@ -56,8 +57,8 @@ func TriggerIngestionTask[T any](d *Client, dataSourceName string, entries []T)
 	return taskID, err
 }
 
-// AwaitTaskCompletion waits for the task to complete. Function timeouts with an error after awaitTimeout nanoseconds.
-func AwaitTaskCompletion(client *Client, taskID string, awaitTimeout time.Duration, tickerDuration time.Duration) error {
+// awaitTaskCompletion waits for the task to complete. Function timeouts with an error after awaitTimeout nanoseconds.
+func awaitTaskCompletion(client *Client, taskID string, awaitTimeout time.Duration, tickerDuration time.Duration) error {
 	ticker := time.NewTicker(tickerDuration)
 	defer ticker.Stop()
 	afterTimeout := time.After(awaitTimeout)
@@ -79,8 +80,8 @@ func AwaitTaskCompletion(client *Client, taskID string, awaitTimeout time.Durati
 	}
 }
 
-// AwaitTaskStatus waits for the druid task status for the maximum of awaitTimeout duration, querying druid task API.
-func AwaitTaskStatus(client *Client, taskID string, status string, awaitTimeout time.Duration, tickerDuration time.Duration) error {
+// awaitTaskStatus waits for the druid task status for the maximum of awaitTimeout duration, querying druid task API.
+func awaitTaskStatus(client *Client, taskID string, status string, awaitTimeout time.Duration, tickerDuration time.Duration) error {
 	ticker := time.NewTicker(tickerDuration)
 	defer ticker.Stop()
 	afterTimeout := time.After(awaitTimeout)
@@ -101,24 +102,24 @@ func AwaitTaskStatus(client *Client, taskID string, status string, awaitTimeout
 	}
 }
 
-// RunInlineIngestionTask initiates inline ingestion task with druid client and runs until it is complete.
-func RunInlineIngestionTask[T any](client *Client, dataSourceName string, entries []T, recordsCount int) error {
-	taskID, err := TriggerIngestionTask(client, dataSourceName, entries)
+// runInlineIngestionTask initiates inline ingestion task with druid client and runs until it is complete.
+func runInlineIngestionTask[T any](client *Client, dataSourceName string, entries []T, recordsCount int) error {
+	taskID, err := triggerIngestionTask(client, dataSourceName, entries)
 	if err != nil {
 		return err
 	}
 
-	err = AwaitTaskCompletion(client, taskID, 180*time.Second, 500*time.Millisecond)
+	err = awaitTaskCompletion(client, taskID, 180*time.Second, 500*time.Millisecond)
 	if err != nil {
 		return err
 	}
 
-	err = client.Metadata(WithMetadataQueryTicker(500*time.Millisecond), WithMetadataQueryTimeout(180*time.Second)).AwaitDataSourceAvailable(dataSourceName)
+	err = client.metadata(WithMetadataQueryTicker(500*time.Millisecond), WithMetadataQueryTimeout(180*time.Second)).awaitDataSourceAvailable(dataSourceName)
 	if err != nil {
 		return err
 	}
 
-	err = client.Metadata().AwaitRecordsCount(dataSourceName, recordsCount)
+	err = client.metadata().awaitRecordsCount(dataSourceName, recordsCount)
 	if err != nil {
 		return err
 	}
@@ -153,7 +154,7 @@ func TestTaskService(t *testing.T) {
 	require.NoError(t, err, "druid services should be up with no error")
 
 	// Test create ingestion task -> get status -> complete sequence.
-	RunInlineIngestionTask(d, "test-submit-task-datasource", testObjects, 2)
+	runInlineIngestionTask(d, "test-submit-task-datasource", testObjects, 2)
 	require.NoError(t, err, "error should be nil")
 }
 
@@ -184,10 +185,10 @@ func TestTerminateTask(t *testing.T) {
 	require.NoError(t, err, "druid services should be up with no error")
 
 	// Test create ingestion task -> get status -> terminate sequence.
-	taskID, err := TriggerIngestionTask(d, "test-terminate-task-datasource", testObjects)
+	taskID, err := triggerIngestionTask(d, "test-terminate-task-datasource", testObjects)
 	require.NoError(t, err, "error should be nil")
 
-	err = AwaitTaskStatus(d, taskID, "RUNNING", 180*time.Second, 200*time.Millisecond)
+	err = awaitTaskStatus(d, taskID, "RUNNING", 180*time.Second, 200*time.Millisecond)
 	require.NoError(t, err, "error should be nil")
 
 	shutdownTaskID, err := d.Tasks().Shutdown(taskID)