data_drain 0.1.19 → 0.2.1

data/docs/glue_pyspark_example.py

@@ -0,0 +1,60 @@
+"""
+Script de AWS Glue (PySpark) compatible con DataDrain::GlueRunner.
+
+Crear el Job en la consola de AWS Glue (Spark 4.0+) y usar este script como base.
+Argumentos requeridos: JOB_NAME, start_date, end_date, s3_bucket, s3_folder,
+db_url, db_user, db_password, db_table, partition_by.
+
+Personalizar la sección de columnas derivadas según las partition_keys de cada tabla.
+"""
+
+import sys
+from awsglue.utils import getResolvedOptions
+from pyspark.context import SparkContext
+from awsglue.context import GlueContext
+from awsglue.job import Job
+from pyspark.sql.functions import col, year, month
+
+args = getResolvedOptions(sys.argv, [
+    'JOB_NAME', 'start_date', 'end_date', 's3_bucket', 's3_folder',
+    'db_url', 'db_user', 'db_password', 'db_table', 'partition_by'
+])
+
+sc = SparkContext()
+glueContext = GlueContext(sc)
+spark = glueContext.spark_session
+job = Job(glueContext)
+job.init(args['JOB_NAME'], args)
+
+options = {
+    "url": args['db_url'],
+    "dbtable": args['db_table'],
+    "user": args['db_user'],
+    "password": args['db_password'],
+    "sampleQuery": (
+        f"SELECT * FROM {args['db_table']} "
+        f"WHERE created_at >= '{args['start_date']}' "
+        f"AND created_at < '{args['end_date']}'"
+    )
+}
+
+df = spark.read.format("jdbc").options(**options).load()
+
+# Agregar columnas derivadas necesarias para las particiones.
+# isp_id ya existe en la tabla fuente — solo agregar las que se calculan.
+# Personalizar esta seccion segun las partition_keys de cada tabla.
+df_final = (
+    df.withColumn("year", year(col("created_at")))
+      .withColumn("month", month(col("created_at")))
+)
+
+output_path = f"s3://{args['s3_bucket']}/{args['s3_folder']}/"
+partitions = args['partition_by'].split(",")
+
+(df_final.write.mode("overwrite")
+    .partitionBy(*partitions)
+    .format("parquet")
+    .option("compression", "zstd")
+    .save(output_path))
+
+job.commit()

data/lib/data_drain/engine.rb

@@ -5,6 +5,7 @@ require "pg"
 
 module DataDrain
   # Motor principal de extracción y purga de datos (DataDrain).
+  # rubocop:disable Metrics/ClassLength, Metrics/AbcSize, Metrics/MethodLength, Naming/AccessorMethodName
   #
   # Orquesta el flujo ETL desde PostgreSQL hacia un Data Lake analítico
   # delegando la interacción del almacenamiento al adaptador configurado.
@@ -21,29 +22,31 @@ module DataDrain
     # @option options [Array<String, Symbol>] :partition_keys Columnas para particionar.
     # @option options [String] :primary_key (Opcional) Clave primaria para borrado. Por defecto 'id'.
     # @option options [String] :where_clause (Opcional) Condición SQL extra.
-    # @option options [Boolean] :skip_export (Opcional) Si es true, no realiza el export a Parquet, solo validación y purga.
+    # @option options [Boolean] :skip_export (Opcional) Si true, no exporta
+    #   a Parquet — solo valida y purga (para uso con GlueRunner).
     def initialize(options)
-      @start_date     = options.fetch(:start_date).beginning_of_day
-      
-      # Usamos el inicio del día siguiente como límite superior estricto (<)
-      # Esto evita problemas de precisión con los microsegundos al usar end_of_day
-      @end_date       = options.fetch(:end_date).to_date.next_day.beginning_of_day
-      
-      @table_name     = options.fetch(:table_name)
-      @folder_name    = options.fetch(:folder_name, @table_name)
-      @select_sql     = options.fetch(:select_sql, "*")
-      @partition_keys = options.fetch(:partition_keys)
-      @primary_key    = options.fetch(:primary_key, "id")
-      @where_clause   = options[:where_clause]
-      @bucket         = options[:bucket]
-      @skip_export    = options.fetch(:skip_export, false)
+      @start_date = options.fetch(:start_date).beginning_of_day
+
+      @end_date = options.fetch(:end_date).to_date.next_day.beginning_of_day
 
-      @config  = DataDrain.configuration
-      @logger  = @config.logger
+      @table_name = options.fetch(:table_name)
+      Validations.validate_identifier!(:table_name, @table_name)
+
+      @folder_name = options.fetch(:folder_name, @table_name)
+      @select_sql = options.fetch(:select_sql, "*")
+      @partition_keys = options.fetch(:partition_keys)
+      @primary_key = options.fetch(:primary_key, "id")
+      Validations.validate_identifier!(:primary_key, @primary_key)
+      @where_clause = options[:where_clause]
+      @bucket = options[:bucket]
+      @skip_export = options.fetch(:skip_export, false)
+
+      @config = DataDrain.configuration
+      @logger = @config.logger
       @adapter = DataDrain::Storage.adapter
 
       database = DuckDB::Database.open(":memory:")
-      @duckdb  = database.connect
+      @duckdb = database.connect
     end
 
     # Ejecuta el flujo completo del motor: Setup, Conteo, Exportación (opcional), Verificación y Purga.
@@ -51,7 +54,8 @@ module DataDrain
     # @return [Boolean] `true` si el proceso finalizó con éxito, `false` si falló la integridad.
     def call
       start_time = Process.clock_gettime(Process::CLOCK_MONOTONIC)
-      safe_log(:info, "engine.start", { table: @table_name, start_date: @start_date.to_date, end_date: @end_date.to_date })
+      safe_log(:info, "engine.start",
+               { table: @table_name, start_date: @start_date.to_date, end_date: @end_date.to_date })
 
       setup_duckdb
 
@@ -62,7 +66,8 @@ module DataDrain
 
       if @pg_count.zero?
         duration = Process.clock_gettime(Process::CLOCK_MONOTONIC) - start_time
-        safe_log(:info, "engine.skip_empty", { table: @table_name, duration_s: duration.round(2), db_query_duration_s: db_query_duration.round(2) })
+        safe_log(:info, "engine.skip_empty",
+                 { table: @table_name, duration_s: duration.round(2), db_query_duration_s: db_query_duration.round(2) })
         return true
       end
 
@@ -90,18 +95,19 @@ module DataDrain
 
         duration = Process.clock_gettime(Process::CLOCK_MONOTONIC) - start_time
         safe_log(:info, "engine.complete", {
-          table: @table_name,
-          duration_s: duration.round(2),
-          db_query_duration_s: db_query_duration.round(2),
-          export_duration_s: export_duration.round(2),
-          integrity_duration_s: integrity_duration.round(2),
-          purge_duration_s: purge_duration.round(2),
-          count: @pg_count
-        })
+                   table: @table_name,
+                   duration_s: duration.round(2),
+                   db_query_duration_s: db_query_duration.round(2),
+                   export_duration_s: export_duration.round(2),
+                   integrity_duration_s: integrity_duration.round(2),
+                   purge_duration_s: purge_duration.round(2),
+                   count: @pg_count
+                 })
         true
       else
         duration = Process.clock_gettime(Process::CLOCK_MONOTONIC) - start_time
-        safe_log(:error, "engine.integrity_error", { table: @table_name, duration_s: duration.round(2), count: @pg_count })
+        safe_log(:error, "engine.integrity_error",
+                 { table: @table_name, duration_s: duration.round(2), count: @pg_count })
         false
       end
     end
@@ -142,7 +148,12 @@ module DataDrain
       @adapter.prepare_export_path(@bucket, @folder_name)
 
       # Determinamos el path base de destino según el adaptador
-      dest_path = @config.storage_mode.to_sym == :s3 ? "s3://#{@bucket}/#{@folder_name}/" : File.join(@bucket, @folder_name, "")
+      dest_path = if @config.storage_mode.to_sym == :s3
+                    "s3://#{@bucket}/#{@folder_name}/"
+                  else
+                    File.join(@bucket,
+                              @folder_name, "")
+                  end
 
       pg_sql = "SELECT #{@select_sql} FROM public.#{@table_name} WHERE #{base_where_sql}"
       pg_sql = pg_sql.gsub("'", "''")
@@ -154,7 +165,7 @@ module DataDrain
         ) TO '#{dest_path}'
         (
           FORMAT PARQUET,
-          PARTITION_BY (#{@partition_keys.join(', ')}),
+          PARTITION_BY (#{@partition_keys.join(", ")}),
           COMPRESSION 'ZSTD',
           OVERWRITE_OR_IGNORE 1
         );
@@ -180,7 +191,8 @@ module DataDrain
         return false
       end
 
-      safe_log(:info, "engine.integrity_check", { table: @table_name, pg_count: @pg_count, parquet_count: parquet_result })
+      safe_log(:info, "engine.integrity_check",
+               { table: @table_name, pg_count: @pg_count, parquet_count: parquet_result })
       @pg_count == parquet_result
     end
 
@@ -189,11 +201,11 @@ module DataDrain
       safe_log(:info, "engine.purge_start", { table: @table_name, batch_size: @config.batch_size })
 
       conn = PG.connect(
-        host:     @config.db_host,
-        port:     @config.db_port,
-        user:     @config.db_user,
+        host: @config.db_host,
+        port: @config.db_port,
+        user: @config.db_user,
         password: @config.db_pass,
-        dbname:   @config.db_name
+        dbname: @config.db_name
       )
 
       unless @config.idle_in_transaction_session_timeout.nil?
@@ -223,10 +235,10 @@ module DataDrain
         # Heartbeat cada 100 lotes para monitorear procesos largos de 1TB
         if (batches_processed % 100).zero?
           safe_log(:info, "engine.purge_heartbeat", {
-            table: @table_name,
-            batches_processed_count: batches_processed,
-            rows_deleted_count: total_deleted
-          })
+                     table: @table_name,
+                     batches_processed_count: batches_processed,
+                     rows_deleted_count: total_deleted
+                   })
         end
 
         sleep(@config.throttle_delay) if @config.throttle_delay.positive?
@@ -235,4 +247,5 @@ module DataDrain
       conn&.close
     end
   end
+  # rubocop:enable Metrics/ClassLength, Metrics/AbcSize, Metrics/MethodLength, Naming/AccessorMethodName
 end

data/lib/data_drain/file_ingestor.rb

@@ -6,6 +6,8 @@ module DataDrain
   # aplicando compresión ZSTD y particionamiento Hive.
   class FileIngestor
     include Observability
+    # rubocop:disable Metrics/AbcSize, Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity,
+    #   Metrics/MethodLength
 
     # @param options [Hash] Opciones de ingestión.
     # @option options [String] :source_path Ruta absoluta al archivo local.
@@ -14,19 +16,20 @@ module DataDrain
     # @option options [String] :select_sql (Opcional) Sentencia SELECT para transformar datos al vuelo.
     # @option options [Boolean] :delete_after_upload (Opcional) Borra el archivo local al terminar. Por defecto true.
     def initialize(options)
-      @source_path         = options.fetch(:source_path)
-      @folder_name         = options.fetch(:folder_name)
-      @partition_keys      = options.fetch(:partition_keys, [])
-      @select_sql          = options.fetch(:select_sql, "*")
+      @source_path = options.fetch(:source_path)
+      @folder_name = options.fetch(:folder_name)
+      Validations.validate_identifier!(:folder_name, @folder_name)
+      @partition_keys = options.fetch(:partition_keys, [])
+      @select_sql = options.fetch(:select_sql, "*")
       @delete_after_upload = options.fetch(:delete_after_upload, true)
-      @bucket              = options[:bucket]
+      @bucket = options[:bucket]
 
-      @config  = DataDrain.configuration
-      @logger  = @config.logger
+      @config = DataDrain.configuration
+      @logger = @config.logger
       @adapter = DataDrain::Storage.adapter
 
       database = DuckDB::Database.open(":memory:")
-      @duckdb  = database.connect
+      @duckdb = database.connect
     end
 
     # Ejecuta el flujo de ingestión.
@@ -52,7 +55,11 @@ module DataDrain
       step_start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
       source_count = @duckdb.query("SELECT COUNT(*) FROM #{reader_function}").first.first
       source_query_duration = Process.clock_gettime(Process::CLOCK_MONOTONIC) - step_start
-      safe_log(:info, "file_ingestor.count", { source_path: @source_path, count: source_count, source_query_duration_s: source_query_duration.round(2) })
+      safe_log(:info, "file_ingestor.count", {
+                 source_path: @source_path,
+                 count: source_count,
+                 source_query_duration_s: source_query_duration.round(2)
+               })
 
       if source_count.zero?
         cleanup_local_file
@@ -63,9 +70,14 @@ module DataDrain
 
       # 2. Exportación / Subida
       @adapter.prepare_export_path(@bucket, @folder_name)
-      dest_path = @config.storage_mode.to_sym == :s3 ? "s3://#{@bucket}/#{@folder_name}/" : File.join(@bucket, @folder_name, "")
+      dest_path = if @config.storage_mode.to_sym == :s3
+                    "s3://#{@bucket}/#{@folder_name}/"
+                  else
+                    File.join(@bucket,
+                              @folder_name, "")
+                  end
 
-      partition_clause = @partition_keys.any? ? "PARTITION_BY (#{@partition_keys.join(', ')})," : ""
+      partition_clause = @partition_keys.any? ? "PARTITION_BY (#{@partition_keys.join(", ")})," : ""
 
       query = <<~SQL
         COPY (
@@ -87,18 +99,19 @@ module DataDrain
 
       duration = Process.clock_gettime(Process::CLOCK_MONOTONIC) - start_time
       safe_log(:info, "file_ingestor.complete", {
-        source_path: @source_path,
-        duration_s: duration.round(2),
-        source_query_duration_s: source_query_duration.round(2),
-        export_duration_s: export_duration.round(2),
-        count: source_count
-      })
+                 source_path: @source_path,
+                 duration_s: duration.round(2),
+                 source_query_duration_s: source_query_duration.round(2),
+                 export_duration_s: export_duration.round(2),
+                 count: source_count
+               })
 
       cleanup_local_file
       true
     rescue DuckDB::Error => e
       duration = Process.clock_gettime(Process::CLOCK_MONOTONIC) - start_time
-      safe_log(:error, "file_ingestor.duckdb_error", { source_path: @source_path }.merge(exception_metadata(e)).merge(duration_s: duration.round(2)))
+      safe_log(:error, "file_ingestor.duckdb_error",
+               { source_path: @source_path }.merge(exception_metadata(e)).merge(duration_s: duration.round(2)))
       false
     ensure
       @duckdb&.close
@@ -109,11 +122,11 @@ module DataDrain
     # @api private
     def determine_reader
       case File.extname(@source_path).downcase
-      when '.csv'
+      when ".csv"
         "read_csv_auto('#{@source_path}')"
-      when '.json'
+      when ".json"
         "read_json_auto('#{@source_path}')"
-      when '.parquet'
+      when ".parquet"
         "read_parquet('#{@source_path}')"
       else
         raise DataDrain::Error, "Formato de archivo no soportado para ingestión: #{@source_path}"
@@ -122,10 +135,12 @@ module DataDrain
 
     # @api private
     def cleanup_local_file
-      if @delete_after_upload && File.exist?(@source_path)
-        File.delete(@source_path)
-        safe_log(:info, "file_ingestor.cleanup", { source_path: @source_path })
-      end
+      return unless @delete_after_upload && File.exist?(@source_path)
+
+      File.delete(@source_path)
+      safe_log(:info, "file_ingestor.cleanup", { source_path: @source_path })
     end
   end
+  # rubocop:enable Metrics/AbcSize, Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity,
+  #   Metrics/MethodLength
 end

data/lib/data_drain/record.rb

@@ -24,10 +24,28 @@ module DataDrain
     class_attribute :folder_name
     class_attribute :partition_keys
 
+    # Cierra la conexión DuckDB del thread actual y limpia Thread.current.
+    # Idempotente: llamarlo varias veces no levanta.
+    #
+    # Útil en middlewares de Sidekiq/Puma para evitar memory leak en threads
+    # de larga vida.
+    #
+    # @return [void]
+    def self.disconnect!
+      entry = Thread.current[:data_drain_duckdb]
+      Thread.current[:data_drain_duckdb] = nil
+      return unless entry
+
+      entry[:conn]&.close
+      entry[:db]&.close
+    rescue StandardError # rubocop:disable Lint/SuppressedException
+    end
+
     # Retorna la conexión persistente a DuckDB en memoria para el hilo (Thread) actual.
     # Esto previene tener que recargar extensiones (como httpfs) en cada consulta.
     #
     # @return [DuckDB::Connection] Conexión activa a DuckDB.
+    # rubocop:disable Metrics/AbcSize
     def self.connection
       Thread.current[:data_drain_duckdb] ||= begin
         db = DuckDB::Database.open(":memory:")
@@ -42,6 +60,7 @@ module DataDrain
       end
       Thread.current[:data_drain_duckdb][:conn]
     end
+    # rubocop:enable Metrics/AbcSize
 
     # Consulta registros en el Data Lake filtrando por claves de partición.
     #
@@ -52,7 +71,7 @@ module DataDrain
       path = build_query_path(partitions)
 
       sql = <<~SQL
-        SELECT #{attribute_names.join(', ')}
+        SELECT #{attribute_names.join(", ")}
         FROM read_parquet('#{path}')
         ORDER BY created_at DESC
         LIMIT #{limit}
@@ -73,7 +92,7 @@ module DataDrain
       safe_id = id.to_s.gsub("'", "''")
 
       sql = <<~SQL
-        SELECT #{attribute_names.join(', ')}
+        SELECT #{attribute_names.join(", ")}
         FROM read_parquet('#{path}')
         WHERE id = '#{safe_id}'
         LIMIT 1
@@ -97,7 +116,7 @@ module DataDrain
     # @return [String] Representación legible en consola.
     def inspect
       inspection = attributes.map do |name, value|
-        "#{name}: #{value.nil? ? 'nil' : value.inspect}"
+        "#{name}: #{value.nil? ? "nil" : value.inspect}"
       end.compact.join(", ")
 
       "#<#{self.class} #{inspection}>"
@@ -118,6 +137,7 @@ module DataDrain
       # @param sql [String]
       # @param columns [Array<String>]
       # @return [Array<DataDrain::Record>]
+      # rubocop:disable Metrics/MethodLength
       def execute_and_instantiate(sql, columns)
         @logger = DataDrain.configuration.logger
         begin
@@ -133,5 +153,6 @@ module DataDrain
         end
       end
     end
+    # rubocop:enable Metrics/MethodLength
   end
 end

data/lib/data_drain/storage/s3.rb

@@ -4,21 +4,59 @@ module DataDrain
   module Storage
     # Implementación del adaptador de almacenamiento para Amazon S3.
     class S3 < Base
+      # rubocop:disable Metrics/AbcSize, Metrics/CyclomaticComplexity, Metrics/MethodLength
       # Carga la extensión httpfs en DuckDB e inyecta las credenciales de AWS.
+      # Si aws_access_key_id y aws_secret_access_key están seteados, usa
+      # credenciales explícitas. Si no, usa credential_chain (IAM role, env vars,
+      # ~/.aws/credentials).
       # @param connection [DuckDB::Connection]
+      # @raise [DataDrain::ConfigurationError] si aws_region no está configurado
       def setup_duckdb(connection)
         connection.query("INSTALL httpfs; LOAD httpfs;")
-        connection.query("SET s3_region='#{@config.aws_region}';")
-        connection.query("SET s3_access_key_id='#{@config.aws_access_key_id}';")
-        connection.query("SET s3_secret_access_key='#{@config.aws_secret_access_key}';")
+        create_s3_secret(connection)
       end
 
+      private
+
+      # @param connection [DuckDB::Connection]
+      # @raise [DataDrain::ConfigurationError]
+      def create_s3_secret(connection)
+        region = @config.aws_region
+        raise DataDrain::ConfigurationError, "aws_region es obligatorio para storage_mode=:s3" if region.nil?
+
+        if @config.aws_access_key_id && @config.aws_secret_access_key
+          connection.query(<<~SQL)
+            CREATE OR REPLACE SECRET data_drain_s3 (
+              TYPE S3,
+              KEY_ID '#{escape_sql(@config.aws_access_key_id)}',
+              SECRET '#{escape_sql(@config.aws_secret_access_key)}',
+              REGION '#{escape_sql(region)}'
+            );
+          SQL
+        else
+          connection.query(<<~SQL)
+            CREATE OR REPLACE SECRET data_drain_s3 (
+              TYPE S3,
+              PROVIDER credential_chain,
+              REGION '#{escape_sql(region)}'
+            );
+          SQL
+        end
+      end
+
+      # @param value [String]
+      # @return [String]
+      def escape_sql(value)
+        value.to_s.gsub("'", "''")
+      end
+
+      public
+
       # @param bucket [String]
       # @param folder_name [String]
       # @param partition_path [String, nil]
       # @return [String]
       def build_path(bucket, folder_name, partition_path)
-        # En S3, el base_path actúa como el nombre del bucket
         base = File.join(bucket, folder_name)
         base = File.join(base, partition_path) if partition_path && !partition_path.empty?
         "s3://#{base}/**/*.parquet"
@@ -40,7 +78,7 @@ module DataDrain
           val = partitions[key]
           val.nil? || val.to_s.empty? ? "#{key}=[^/]+" : "#{key}=#{val}"
         end
-        pattern_regex = Regexp.new("^#{folder_name}/#{regex_parts.join('/')}")
+        pattern_regex = Regexp.new("^#{folder_name}/#{regex_parts.join("/")}")
 
         objects_to_delete = []
         prefix = "#{folder_name}/"
@@ -58,7 +96,10 @@ module DataDrain
 
       private
 
-      # @api private
+      # @param client [Aws::S3::Client]
+      # @param bucket [String]
+      # @param objects_to_delete [Array<Hash>]
+      # @return [Integer]
       def delete_in_batches(client, bucket, objects_to_delete)
         return 0 if objects_to_delete.empty?
 
@@ -70,5 +111,6 @@ module DataDrain
         deleted_count
       end
     end
+    # rubocop:enable Metrics/AbcSize, Metrics/CyclomaticComplexity, Metrics/MethodLength
   end
 end

data/lib/data_drain/validations.rb

@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+module DataDrain
+  # Módulo de validación de configuración para prevenir errores de uso.
+  module Validations
+    IDENTIFIER_REGEX = /\A[a-zA-Z_][a-zA-Z0-9_]*\z/
+
+    module_function
+
+    def validate_identifier!(name, value)
+      return if IDENTIFIER_REGEX.match?(value.to_s)
+
+      raise DataDrain::ConfigurationError,
+            "#{name} '#{value}' no es un identificador SQL válido"
+    end
+  end
+end

data/lib/data_drain/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module DataDrain
-  VERSION = "0.1.19"
+  VERSION = "0.2.1"
 end

data/lib/data_drain.rb

@@ -4,6 +4,7 @@ require "active_model"
 require_relative "data_drain/version"
 require_relative "data_drain/errors"
 require_relative "data_drain/configuration"
+require_relative "data_drain/validations"
 require_relative "data_drain/storage"
 require_relative "data_drain/observability"
 require_relative "data_drain/engine"
@@ -15,6 +16,7 @@ require_relative "data_drain/glue_runner"
 require_relative "data_drain/types/json_type"
 ActiveModel::Type.register(:json, DataDrain::Types::JsonType)
 
+# DSL para extraer, archivar y purgar datos entre PostgreSQL y un Data Lake en Parquet.
 module DataDrain
   class << self
     # @return [DataDrain::Configuration]