sqliteexporter/factory.go

// Copyright 2024 William Perron. All rights reserved. MIT License.
package sqliteexporter

import (
	"context"
	"database/sql"
	_ "embed"
	"fmt"

	"go.opentelemetry.io/collector/component"
	"go.opentelemetry.io/collector/consumer"
	"go.opentelemetry.io/collector/exporter"
	"go.opentelemetry.io/collector/exporter/exporterhelper"
	sdktrace "go.opentelemetry.io/otel/sdk/trace"

	"go.wperron.io/sqliteexporter/internal/metadata"
)

//go:embed migrations/20240120195122_init.up.sql
var initScript string

func NewFactory() exporter.Factory {
	return exporter.NewFactory(
		metadata.Type,
		createDefaultConfig,
		exporter.WithTraces(createTracesExporter, metadata.TracesStability),
		exporter.WithMetrics(nil, metadata.MetricsStability),
		exporter.WithLogs(nil, metadata.LogsStability),
	)
}

func createDefaultConfig() component.Config {
	return &Config{}
}

func createTracesExporter(
	ctx context.Context,
	set exporter.CreateSettings,
	cfg component.Config,
) (exporter.Traces, error) {
	conf := cfg.(*Config)

	se, err := newSqliteExporter(conf)
	if err != nil {
		return nil, fmt.Errorf("failed to create sqlite exporter: %w", err)
	}

	return exporterhelper.NewTracesExporter(
		ctx, set, cfg,
		se.ConsumeTraces,
		exporterhelper.WithStart(se.Start),
		exporterhelper.WithShutdown(se.Shutdown),
		exporterhelper.WithCapabilities(consumer.Capabilities{MutatesData: false}),
	)
}

func newSqliteExporter(cfg *Config) (*sqliteExporter, error) {
	db, err := sql.Open("sqlite3", cfg.Path)
	if err != nil {
		return nil, fmt.Errorf("couldn't open sqlite3 database: %w", err)
	}

	// IMPORTANT: database/sql opens a connection pool by default, but sqlite
	// only allows a single connection to be open at the same time.
	db.SetMaxOpenConns(1)

	tx, _ := db.Begin()
	defer tx.Commit()

	if _, err := tx.Exec(initScript); err != nil {
		return nil, fmt.Errorf("failed to run migrations: %w", err)
	}

	return &sqliteExporter{
		db: db,
	}, nil
}

func NewSqliteSDKTraceExporter(cfg *Config) (sdktrace.SpanExporter, error) {
	return newSqliteExporter(cfg)
}

func NewSqliteSDKTraceExporterWithDB(db *sql.DB) (sdktrace.SpanExporter, error) {
	tx, _ := db.Begin()
	defer tx.Commit()

	if _, err := tx.Exec(initScript); err != nil {
		return nil, fmt.Errorf("failed to run migrations: %w", err)
	}

	return &sqliteExporter{
		db: db,
	}, nil
}